From 16d691332de1449e22f4a2be1ae2a588142fc212 Mon Sep 17 00:00:00 2001 From: cc Date: Sun, 28 Jun 2026 21:50:37 -0700 Subject: [PATCH 01/29] spec: planning group --- .../control-docs-housekeeping/.openspec.yaml | 2 ++ .../control-docs-housekeeping/README.md | 3 ++ .../control-docs-housekeeping/design.md | 26 +++++++++++++++++ .../control-docs-housekeeping/proposal.md | 25 ++++++++++++++++ .../specs/control-docs-housekeeping/spec.md | 15 ++++++++++ .../control-docs-housekeeping/tasks.md | 17 +++++++++++ .../changes/group-aware-ik-rrt/.openspec.yaml | 2 ++ openspec/changes/group-aware-ik-rrt/README.md | 3 ++ openspec/changes/group-aware-ik-rrt/design.md | 27 +++++++++++++++++ .../changes/group-aware-ik-rrt/proposal.md | 24 +++++++++++++++ .../specs/group-aware-ik-rrt/spec.md | 22 ++++++++++++++ openspec/changes/group-aware-ik-rrt/tasks.md | 19 ++++++++++++ .../.openspec.yaml | 2 ++ .../manipulation-module-group-api/README.md | 3 ++ .../manipulation-module-group-api/design.md | 28 ++++++++++++++++++ .../manipulation-module-group-api/proposal.md | 25 ++++++++++++++++ .../manipulation-module-group-api/spec.md | 23 +++++++++++++++ .../manipulation-module-group-api/tasks.md | 18 ++++++++++++ .../planning-groups-foundation/.openspec.yaml | 2 ++ .../planning-groups-foundation/README.md | 3 ++ .../planning-groups-foundation/design.md | 29 +++++++++++++++++++ .../planning-groups-foundation/proposal.md | 28 ++++++++++++++++++ .../manipulation-planning-config/spec.md | 8 +++++ .../specs/planning-group-foundation/spec.md | 22 ++++++++++++++ .../planning-groups-foundation/tasks.md | 18 ++++++++++++ .../.openspec.yaml | 2 ++ .../planning-world-monitor-groups/README.md | 3 ++ .../planning-world-monitor-groups/design.md | 27 +++++++++++++++++ .../planning-world-monitor-groups/proposal.md | 26 +++++++++++++++++ .../planning-world-monitor-groups/spec.md | 22 ++++++++++++++ .../planning-world-monitor-groups/tasks.md | 18 ++++++++++++ .../viser-group-planning-ui/.openspec.yaml | 2 ++ .../changes/viser-group-planning-ui/README.md | 3 ++ .../changes/viser-group-planning-ui/design.md | 28 ++++++++++++++++++ .../viser-group-planning-ui/proposal.md | 25 ++++++++++++++++ .../specs/viser-group-planning-ui/spec.md | 22 ++++++++++++++ .../changes/viser-group-planning-ui/tasks.md | 17 +++++++++++ openspec/config.yaml | 20 +++++++++++++ 38 files changed, 609 insertions(+) create mode 100644 openspec/changes/control-docs-housekeeping/.openspec.yaml create mode 100644 openspec/changes/control-docs-housekeeping/README.md create mode 100644 openspec/changes/control-docs-housekeeping/design.md create mode 100644 openspec/changes/control-docs-housekeeping/proposal.md create mode 100644 openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md create mode 100644 openspec/changes/control-docs-housekeeping/tasks.md create mode 100644 openspec/changes/group-aware-ik-rrt/.openspec.yaml create mode 100644 openspec/changes/group-aware-ik-rrt/README.md create mode 100644 openspec/changes/group-aware-ik-rrt/design.md create mode 100644 openspec/changes/group-aware-ik-rrt/proposal.md create mode 100644 openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md create mode 100644 openspec/changes/group-aware-ik-rrt/tasks.md create mode 100644 openspec/changes/manipulation-module-group-api/.openspec.yaml create mode 100644 openspec/changes/manipulation-module-group-api/README.md create mode 100644 openspec/changes/manipulation-module-group-api/design.md create mode 100644 openspec/changes/manipulation-module-group-api/proposal.md create mode 100644 openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md create mode 100644 openspec/changes/manipulation-module-group-api/tasks.md create mode 100644 openspec/changes/planning-groups-foundation/.openspec.yaml create mode 100644 openspec/changes/planning-groups-foundation/README.md create mode 100644 openspec/changes/planning-groups-foundation/design.md create mode 100644 openspec/changes/planning-groups-foundation/proposal.md create mode 100644 openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md create mode 100644 openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md create mode 100644 openspec/changes/planning-groups-foundation/tasks.md create mode 100644 openspec/changes/planning-world-monitor-groups/.openspec.yaml create mode 100644 openspec/changes/planning-world-monitor-groups/README.md create mode 100644 openspec/changes/planning-world-monitor-groups/design.md create mode 100644 openspec/changes/planning-world-monitor-groups/proposal.md create mode 100644 openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md create mode 100644 openspec/changes/planning-world-monitor-groups/tasks.md create mode 100644 openspec/changes/viser-group-planning-ui/.openspec.yaml create mode 100644 openspec/changes/viser-group-planning-ui/README.md create mode 100644 openspec/changes/viser-group-planning-ui/design.md create mode 100644 openspec/changes/viser-group-planning-ui/proposal.md create mode 100644 openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md create mode 100644 openspec/changes/viser-group-planning-ui/tasks.md create mode 100644 openspec/config.yaml diff --git a/openspec/changes/control-docs-housekeeping/.openspec.yaml b/openspec/changes/control-docs-housekeeping/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/control-docs-housekeeping/README.md b/openspec/changes/control-docs-housekeeping/README.md new file mode 100644 index 0000000000..ee333ffc5a --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/README.md @@ -0,0 +1,3 @@ +# control-docs-housekeeping + +Extract PR 6 from reference branch cc/spec/movegroup: peel off control/coordinator/generated/docs/housekeeping changes not central to planning-group semantics. diff --git a/openspec/changes/control-docs-housekeeping/design.md b/openspec/changes/control-docs-housekeeping/design.md new file mode 100644 index 0000000000..c42efe8508 --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/design.md @@ -0,0 +1,26 @@ +## Context + +These files do not obviously belong to the core planning-group review. Some may be genuine integration requirements; others may be incidental edits from the large branch. + +## Goals / Non-Goals + +**Goals:** +- Separate orthogonal changes from the planning-group PR stack. +- Keep generated files and docs synchronized only when required. +- Provide a clear report of accidental or questionable files. + +**Non-Goals:** +- Do not change planning-group semantics. +- Do not include Viser UI or planner algorithm changes. +- Do not include files just because they exist on the reference branch. + +## Decisions + +- Treat this change as a triage/extraction PR, not a mandatory feature PR. +- Include control/task changes only if validation shows they are required or independently useful. +- Keep generated registry changes next to the PR that introduces the corresponding blueprint, unless generation must happen at the end of the stack. + +## Risks / Trade-offs + +- If control changes depend on module API changes, this PR should stack after PR 4. +- If docs describe earlier PR behavior, those docs should move into the earlier PRs instead of this housekeeping PR. diff --git a/openspec/changes/control-docs-housekeeping/proposal.md b/openspec/changes/control-docs-housekeeping/proposal.md new file mode 100644 index 0000000000..436f94fb41 --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/proposal.md @@ -0,0 +1,25 @@ +## Why + +The current branch contains small control, generated-registry, docs, and housekeeping changes mixed into the planning-group refactor. These should be peeled off so reviewers can decide whether they are required or accidental. + +## What Changes + +- Extract control/coordinator task and tick-loop changes if they are truly required. +- Extract generated blueprint registry updates if they are caused by earlier PRs. +- Extract remaining docs/readme cleanup and packaging/config touches. +- Report any file that appears accidental rather than including it blindly. + +## Capabilities + +### New Capabilities +- `control-docs-housekeeping`: Distribution plan for orthogonal control, generated, docs, and repo housekeeping changes. + +### Modified Capabilities +- `control-coordinator-integration`: Control/coordinator behavior updates only if validated as required by the planning-group stack. + +## Impact + +- Base branch: preferably PR 4 `manipulation-module-group-api`, or `main` for truly independent control/docs changes. +- Reference implementation: `cc/spec/movegroup`. +- Primary candidate files: `dimos/control/*`, `dimos/e2e_tests/test_control_coordinator.py`, generated blueprint files, docs/readmes, `pyproject.toml`, `AGENTS.md`. +- Out of scope: planning-group core, planners, module API, and Viser UI. diff --git a/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md b/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md new file mode 100644 index 0000000000..cc200026e1 --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md @@ -0,0 +1,15 @@ +## ADDED Requirements + +### Requirement: Orthogonal changes must be reviewed separately +Changes that are not necessary for planning-group semantics MUST be extracted into a separate review or explicitly excluded from the stack. + +#### Scenario: Candidate file is unrelated to planning groups +- **WHEN** a file change is docs-only, generated-only, control-only, or repo-housekeeping-only +- **THEN** it is reviewed in this housekeeping change or reported as intentionally excluded + +### Requirement: Generated files must match their source changes +Generated registry files MUST be included only with the source blueprint/config change that requires regeneration, or in a final generated-files PR. + +#### Scenario: Blueprint registry changes are generated +- **WHEN** a generated registry file differs from `main` +- **THEN** the extraction task identifies the source change and validates regeneration diff --git a/openspec/changes/control-docs-housekeeping/tasks.md b/openspec/changes/control-docs-housekeeping/tasks.md new file mode 100644 index 0000000000..6831183321 --- /dev/null +++ b/openspec/changes/control-docs-housekeeping/tasks.md @@ -0,0 +1,17 @@ +## 1. Triage candidate files + +- [ ] 1.1 Compare each candidate file against the reference branch and decide whether it is required, independent, generated, docs-only, or accidental. +- [ ] 1.2 Extract only required or intentionally independent files. +- [ ] 1.3 Move docs back into earlier PRs when they explain earlier PR behavior. + +## 2. Control/generated/docs extraction + +- [ ] 2.1 Extract control/task changes if they are intentional. +- [ ] 2.2 Extract generated blueprint registry changes only when their source blueprint is included. +- [ ] 2.3 Extract `pyproject.toml` and `AGENTS.md` only if justified; otherwise report as excluded. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/control/test_control.py dimos/e2e_tests/test_control_coordinator.py dimos/robot/test_all_blueprints.py -q` for included files. +- [ ] 3.2 Run any docs or generated-registry checks required by included changes. +- [ ] 3.3 Return a file-by-file inclusion/exclusion summary. diff --git a/openspec/changes/group-aware-ik-rrt/.openspec.yaml b/openspec/changes/group-aware-ik-rrt/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/group-aware-ik-rrt/README.md b/openspec/changes/group-aware-ik-rrt/README.md new file mode 100644 index 0000000000..872d917e36 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/README.md @@ -0,0 +1,3 @@ +# group-aware-ik-rrt + +Extract PR 3 from reference branch cc/spec/movegroup: make IK and RRT consume explicit planning groups and group-local targets. diff --git a/openspec/changes/group-aware-ik-rrt/design.md b/openspec/changes/group-aware-ik-rrt/design.md new file mode 100644 index 0000000000..6187e24bb9 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/design.md @@ -0,0 +1,27 @@ +## Context + +IK and RRT previously operated mostly through robot-scoped assumptions. With groups, each target must map to an explicit group, local joint order, and target frame. + +## Goals / Non-Goals + +**Goals:** +- Make IK target frames come from the requested planning group. +- Make RRT accept and return group-local joint targets and paths. +- Keep collision checks and global robot-state projection correct. +- Fail clearly when a requested group lacks required pose target metadata. + +**Non-Goals:** +- Do not change `ManipulationModule` public APIs in this PR. +- Do not include visualization changes. +- Do not change control tasks. + +## Decisions + +- Prefer explicit group IDs for all new solver/planner paths. +- Robot-scoped compatibility may resolve through a unique pose-targetable group, but ambiguous robots must fail clearly. +- PinkIK and Viser-style base-pose transforms must validate that robot-scoped `base_link` is compatible with model-root assumptions. + +## Risks / Trade-offs + +- Algorithm tests can become broad quickly. Keep this PR focused on solver/planner contracts and leave module-level behavior to PR 4. +- Some optional dependencies may be unavailable locally; fake dependency tests should stay hermetic. diff --git a/openspec/changes/group-aware-ik-rrt/proposal.md b/openspec/changes/group-aware-ik-rrt/proposal.md new file mode 100644 index 0000000000..0ea8cb9a14 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/proposal.md @@ -0,0 +1,24 @@ +## Why + +After world backends understand planning groups, the solver and planner layer must consume group IDs and group-local targets. This PR keeps algorithm review separate from both backend plumbing and public module/UI migration. + +## What Changes + +- Update PinkIK, Jacobian IK, and Drake optimization IK to resolve target frames from planning groups. +- Update RRT planning to handle group-local joint targets, group candidate selection, collision checking, and result projection. +- Add focused IK/planner tests. + +## Capabilities + +### New Capabilities +- `group-aware-ik-rrt`: IK and RRT planning support explicit planning groups and group-local targets. + +### Modified Capabilities +- `manipulation-planning-algorithms`: Planners and IK solvers no longer rely on robot-scoped end-effector metadata. + +## Impact + +- Base branch: PR 2 `planning-world-monitor-groups`. +- Reference implementation: `cc/spec/movegroup`. +- Primary files: `dimos/manipulation/planning/kinematics/*`, `dimos/manipulation/planning/planners/rrt_planner.py`, related tests. +- Out of scope: public `ManipulationModule` API migration and Viser UI. diff --git a/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md b/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md new file mode 100644 index 0000000000..8d477a3035 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md @@ -0,0 +1,22 @@ +## ADDED Requirements + +### Requirement: IK solvers must use planning-group target frames +IK solvers MUST resolve pose target frames from the requested planning group's `tip_link` rather than from robot-scoped end-effector metadata. + +#### Scenario: Pose IK targets a manipulator group +- **WHEN** a pose target is submitted for a group with `tip_link="tcp"` +- **THEN** the IK solver constrains the `tcp` frame + +### Requirement: Planners must preserve group-local joint ordering +Planners MUST accept and return joint targets in the requested group's local joint order while projecting through full robot state for collision checks. + +#### Scenario: Group joint target uses subset order +- **WHEN** a group target names a subset of robot joints +- **THEN** planning uses the correct full robot state and returns a path scoped to the group + +### Requirement: Algorithms must fail clearly for non-pose-targetable groups +Algorithms that require a pose target frame MUST return an explicit failure when the requested group has no `tip_link`. + +#### Scenario: Pose IK targets a joint-only group +- **WHEN** a pose target is submitted for a group without `tip_link` +- **THEN** the solver returns a no-solution or unsupported result with an explanatory message diff --git a/openspec/changes/group-aware-ik-rrt/tasks.md b/openspec/changes/group-aware-ik-rrt/tasks.md new file mode 100644 index 0000000000..3a13b0a9f7 --- /dev/null +++ b/openspec/changes/group-aware-ik-rrt/tasks.md @@ -0,0 +1,19 @@ +## 1. Extract algorithm changes + +- [ ] 1.1 Start from PR 2 and use `cc/spec/movegroup` as reference. +- [ ] 1.2 Extract PinkIK group target-frame and base-link validation changes. +- [ ] 1.3 Extract Jacobian IK group selection changes. +- [ ] 1.4 Extract Drake optimization IK group target-frame changes. +- [ ] 1.5 Extract RRT planner group-local target and projection changes. + +## 2. Tests + +- [ ] 2.1 Bring over PinkIK tests relevant to group target frames and base-link validation. +- [ ] 2.2 Bring over Jacobian IK selection tests. +- [ ] 2.3 Bring over RRT planner selection/group tests. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/kinematics/test_pink_ik.py dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py dimos/manipulation/planning/planners/test_rrt_planner_selection.py -q`. +- [ ] 3.2 Run targeted mypy on changed algorithm files. +- [ ] 3.3 Optional manual smoke: solve IK and plan a small reachable group pose; print status and path length. diff --git a/openspec/changes/manipulation-module-group-api/.openspec.yaml b/openspec/changes/manipulation-module-group-api/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/manipulation-module-group-api/README.md b/openspec/changes/manipulation-module-group-api/README.md new file mode 100644 index 0000000000..a4cb1c316f --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/README.md @@ -0,0 +1,3 @@ +# manipulation-module-group-api + +Extract PR 4 from reference branch cc/spec/movegroup: expose group-aware planning through ManipulationModule while preserving safe robot-scoped compatibility behavior. diff --git a/openspec/changes/manipulation-module-group-api/design.md b/openspec/changes/manipulation-module-group-api/design.md new file mode 100644 index 0000000000..5e84d72704 --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/design.md @@ -0,0 +1,28 @@ +## Context + +The manipulation module is where planning groups become user-facing. It must expose explicit group APIs while retaining safe behavior for existing robot-scoped wrappers. + +## Goals / Non-Goals + +**Goals:** +- Support explicit group IDs for joint targets, pose targets, IK, previews, and robot info. +- Keep robot-scoped wrappers deterministic: no unique pose group means no successful robot-scoped pose action. +- Preserve safe return contracts for RPC-like public methods. +- Cover compatibility behavior in unit tests. + +**Non-Goals:** +- Do not implement Viser panel behavior in this PR. +- Do not add new planner algorithms. +- Do not change control task behavior unless required by module API compilation. + +## Decisions + +- `get_ee_pose` returns `None` when no unique pose-targetable group exists. +- `plan_to_pose` returns `False` when no unique pose-targetable group exists. +- `inverse_kinematics_single` returns `IKResult(NO_SOLUTION, message=...)` when no unique pose-targetable group exists. +- Explicit group APIs are preferred for multi-group robots. + +## Risks / Trade-offs + +- This file has a large diff. Keep the PR focused on API behavior and tests; do not include UI state-machine changes. +- Some compatibility wrappers are intentionally conservative to avoid silently selecting the wrong group. diff --git a/openspec/changes/manipulation-module-group-api/proposal.md b/openspec/changes/manipulation-module-group-api/proposal.md new file mode 100644 index 0000000000..a74f5a3adf --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/proposal.md @@ -0,0 +1,25 @@ +## Why + +Once engines and algorithms are group-aware, the public manipulation module must expose those capabilities safely. This PR is the reviewer-facing API migration and should stay separate from the Viser UI rewrite. + +## What Changes + +- Add or migrate `ManipulationModule` APIs for group-aware joint targets, pose targets, previews, IK, and robot info. +- Preserve robot-scoped compatibility wrappers with explicit behavior when there is no unique pose-targetable group. +- Update coordinator client and example client call sites. +- Add module/unit/e2e tests for public behavior. + +## Capabilities + +### New Capabilities +- `manipulation-module-group-api`: Public manipulation APIs support planning-group IDs and safe compatibility wrappers. + +### Modified Capabilities +- `manipulation-module-planning`: Robot-scoped pose wrappers no longer assume an arbitrary or hidden end-effector. + +## Impact + +- Base branch: PR 3 `group-aware-ik-rrt`. +- Reference implementation: `cc/spec/movegroup`, including Greptile follow-up behavior for robot-scoped wrappers. +- Primary files: `dimos/manipulation/manipulation_module.py`, coordinator/example clients, manipulation module tests, and planning-group e2e test. +- Out of scope: Viser UI and control/task changes. diff --git a/openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md b/openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md new file mode 100644 index 0000000000..16aa9bb6bc --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md @@ -0,0 +1,23 @@ +## ADDED Requirements + +### Requirement: Manipulation module must expose explicit group planning APIs +The manipulation module MUST allow callers to target planning groups explicitly for joint and pose planning operations. + +#### Scenario: Caller plans to a group pose +- **WHEN** a caller submits a pose target for a valid group ID +- **THEN** the module plans using that group and returns the planning result through the existing public contract + +### Requirement: Robot-scoped pose wrappers must fail safely without a unique pose group +Robot-scoped pose wrapper APIs MUST fail safely when a selected robot has zero or multiple pose-targetable groups. + +#### Scenario: End-effector pose is requested with no unique pose group +- **WHEN** `get_ee_pose` cannot resolve exactly one pose-targetable group +- **THEN** it returns `None` + +#### Scenario: Robot-scoped pose planning is requested with no unique pose group +- **WHEN** `plan_to_pose` cannot resolve exactly one pose-targetable group +- **THEN** it returns `False` + +#### Scenario: Robot-scoped IK is requested with no unique pose group +- **WHEN** `inverse_kinematics_single` cannot resolve exactly one pose-targetable group +- **THEN** it returns an `IKResult` with `NO_SOLUTION` diff --git a/openspec/changes/manipulation-module-group-api/tasks.md b/openspec/changes/manipulation-module-group-api/tasks.md new file mode 100644 index 0000000000..7f50e78293 --- /dev/null +++ b/openspec/changes/manipulation-module-group-api/tasks.md @@ -0,0 +1,18 @@ +## 1. Extract module API changes + +- [ ] 1.1 Start from PR 3 and use `cc/spec/movegroup` as reference. +- [ ] 1.2 Extract group-aware target, IK, preview, robot-info, and execution behavior from `ManipulationModule`. +- [ ] 1.3 Extract coordinator client and example client updates only as needed for the public API. +- [ ] 1.4 Include Greptile follow-up behavior for `get_ee_pose` and `plan_to_pose`. + +## 2. Tests + +- [ ] 2.1 Bring over focused `test_manipulation_unit.py` coverage for group APIs and compatibility wrappers. +- [ ] 2.2 Bring over relevant `test_manipulation_module.py` integration coverage. +- [ ] 2.3 Bring over `dimos/e2e_tests/test_manipulation_planning_groups.py` if it can run against this stacked base. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/manipulation/test_manipulation_unit.py dimos/manipulation/test_manipulation_module.py dimos/e2e_tests/test_manipulation_planning_groups.py -q`. +- [ ] 3.2 Run `uv run mypy dimos/manipulation/manipulation_module.py`. +- [ ] 3.3 Verify no Viser implementation files are included. diff --git a/openspec/changes/planning-groups-foundation/.openspec.yaml b/openspec/changes/planning-groups-foundation/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/planning-groups-foundation/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/planning-groups-foundation/README.md b/openspec/changes/planning-groups-foundation/README.md new file mode 100644 index 0000000000..820f444c27 --- /dev/null +++ b/openspec/changes/planning-groups-foundation/README.md @@ -0,0 +1,3 @@ +# planning-groups-foundation + +Extract PR 1 from reference branch cc/spec/movegroup: introduce planning-group foundations, RobotConfig group generation, manipulator group declarations, and foundation docs/tests. diff --git a/openspec/changes/planning-groups-foundation/design.md b/openspec/changes/planning-groups-foundation/design.md new file mode 100644 index 0000000000..287006f709 --- /dev/null +++ b/openspec/changes/planning-groups-foundation/design.md @@ -0,0 +1,29 @@ +## Context + +The current branch introduces planning groups across the stack, but the reviewable foundation is the abstraction itself: stable group IDs, local/global joint mapping, group definitions, RobotConfig discovery/generation, and removal of robot-scoped end-effector metadata. + +## Goals / Non-Goals + +**Goals:** +- Define planning groups as the only source for chain base links and pose target tip links. +- Preserve robot-scoped `base_link` as placement/weld/strip metadata, not a planning-chain selector. +- Let robot configs create explicit or discovered planning groups. +- Provide small, direct tests that prove group IDs, joint-name mapping, discovery, and strict config validation. + +**Non-Goals:** +- Do not migrate world backends or planners in this PR. +- Do not migrate `ManipulationModule` public APIs in this PR. +- Do not include visualization or control changes. + +## Decisions + +- Use `PlanningGroupDefinition` for config-time group declarations and `PlanningGroup` for resolved runtime group IDs. +- Use globally scoped group IDs such as `/` at API boundaries. +- Keep `RobotModelConfig.base_link` for robot placement semantics only. +- Reject legacy robot-scoped `end_effector_link`; pose target frames come from group `tip_link` or SRDF discovery. +- Keep this PR compatible enough for later stacked PRs; if a later change is needed for tests, prefer a small compatibility wrapper over pulling later files forward. + +## Risks / Trade-offs + +- Earlier PRs may need temporary compatibility paths because the full branch removes legacy behavior across several layers. +- Discovery behavior must be clear: high-level `RobotConfig.to_robot_model_config()` can discover/fallback groups, while direct `RobotModelConfig(...)` requires explicit groups. diff --git a/openspec/changes/planning-groups-foundation/proposal.md b/openspec/changes/planning-groups-foundation/proposal.md new file mode 100644 index 0000000000..c1643bbc5e --- /dev/null +++ b/openspec/changes/planning-groups-foundation/proposal.md @@ -0,0 +1,28 @@ +## Why + +The planning-group refactor is currently bundled in one large branch. Reviewers need a small first PR that establishes the core planning-group abstraction and robot configuration plumbing without also reviewing planner algorithms, module APIs, visualization, or control integration. + +This change extracts PR 1 from the reference branch/worktree `cc/spec/movegroup`. Treat the existing branch as a working reference implementation; do not redesign the feature unless required to keep this slice independently valid. + +## What Changes + +- Add planning-group identifiers, definitions, runtime models, registry helpers, and joint-name conversion utilities. +- Update planning spec/config models so robot-scoped `end_effector_link` is removed and pose target frames are represented by planning-group `tip_link` values. +- Add high-level `RobotConfig` conversion/discovery support for planning groups. +- Update manipulator robot configs to declare or discover explicit planning groups. +- Add foundation-level unit tests and minimal concept docs. + +## Capabilities + +### New Capabilities +- `planning-group-foundation`: Core planning-group data model, discovery, robot config generation, and validation contracts. + +### Modified Capabilities +- `manipulation-planning-config`: Robot model configuration uses planning groups as the source of chain and pose-target metadata. + +## Impact + +- Base branch: `main`. +- Reference implementation: branch/worktree `cc/spec/movegroup`; stabilize to a commit before assigning agents. +- Primary files: `dimos/manipulation/planning/groups/*`, `dimos/manipulation/planning/spec/*`, `dimos/robot/config.py`, `dimos/robot/manipulators/*/config.py`, foundation tests, and planning-group docs. +- Out of scope: world backends, IK/RRT algorithms, `ManipulationModule`, Viser UI, and control changes. diff --git a/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md b/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md new file mode 100644 index 0000000000..ecb76b89b5 --- /dev/null +++ b/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md @@ -0,0 +1,8 @@ +## MODIFIED Requirements + +### Requirement: Robot model configuration must separate placement from planning chains +Robot model configuration MUST treat robot-scoped `base_link` as the link placed by `base_pose` and used by backend weld/strip behavior. Planning-chain base links and pose target frames MUST come from planning groups. + +#### Scenario: Robot has one pose-targetable group +- **WHEN** a robot model has `base_link="base_link"` and a planning group with `base_link="link1"` and `tip_link="tcp"` +- **THEN** placement uses the robot-scoped `base_link` while pose planning targets `tcp` diff --git a/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md b/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md new file mode 100644 index 0000000000..2da6cd9e9e --- /dev/null +++ b/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md @@ -0,0 +1,22 @@ +## ADDED Requirements + +### Requirement: Planning group definitions must identify kinematic chains +The planning configuration MUST represent kinematic chain metadata with planning-group definitions containing group-local joint names, a chain base link, and an optional pose-target tip link. + +#### Scenario: Explicit manipulator group is configured +- **WHEN** a manipulator config declares a group with local joints, `base_link`, and `tip_link` +- **THEN** the resolved robot model exposes that group without relying on a robot-scoped end-effector field + +### Requirement: Robot configuration must reject removed end-effector metadata +Robot configuration MUST NOT accept robot-scoped `end_effector_link` as the source of pose-target metadata. It MUST direct callers to planning-group `tip_link` values or SRDF discovery. + +#### Scenario: Legacy end-effector field is provided +- **WHEN** a robot config is constructed with `end_effector_link` +- **THEN** validation fails with migration guidance + +### Requirement: Group identifiers must be globally scoped +Planning-group APIs MUST use stable group identifiers scoped by robot name so multi-robot stacks can distinguish groups with the same local name. + +#### Scenario: Two robots use manipulator groups +- **WHEN** two robot configs each define a `manipulator` group +- **THEN** their resolved group IDs are distinct diff --git a/openspec/changes/planning-groups-foundation/tasks.md b/openspec/changes/planning-groups-foundation/tasks.md new file mode 100644 index 0000000000..1e9ef2d761 --- /dev/null +++ b/openspec/changes/planning-groups-foundation/tasks.md @@ -0,0 +1,18 @@ +## 1. Extract foundation files + +- [ ] 1.1 Start from `main` and use `cc/spec/movegroup` as the reference implementation. +- [ ] 1.2 Bring over `dimos/manipulation/planning/groups/*`. +- [ ] 1.3 Bring over required planning spec/config changes only. +- [ ] 1.4 Bring over `dimos/robot/config.py` and manipulator config group declarations. + +## 2. Tests and docs + +- [ ] 2.1 Bring over `dimos/manipulation/planning/test_planning_groups.py` and `test_planning_group_utils.py`. +- [ ] 2.2 Bring over `dimos/robot/test_config.py`. +- [ ] 2.3 Bring over only foundation docs needed to explain planning groups and custom-arm config. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/test_planning_groups.py dimos/manipulation/planning/test_planning_group_utils.py dimos/robot/test_config.py -q`. +- [ ] 3.2 Run targeted mypy on changed production files. +- [ ] 3.3 Confirm no files from world/IK/RRT/module/Viser/control were pulled in accidentally. diff --git a/openspec/changes/planning-world-monitor-groups/.openspec.yaml b/openspec/changes/planning-world-monitor-groups/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/planning-world-monitor-groups/README.md b/openspec/changes/planning-world-monitor-groups/README.md new file mode 100644 index 0000000000..c5b27282f4 --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/README.md @@ -0,0 +1,3 @@ +# planning-world-monitor-groups + +Extract PR 2 from reference branch cc/spec/movegroup: teach world backends and world monitors to answer planning-group FK, Jacobian, state, and collision queries. diff --git a/openspec/changes/planning-world-monitor-groups/design.md b/openspec/changes/planning-world-monitor-groups/design.md new file mode 100644 index 0000000000..eb1e36a5e6 --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/design.md @@ -0,0 +1,27 @@ +## Context + +Planning groups become useful when world backends can answer queries for a specific group. This layer should prove that configured groups map correctly onto backend model instances, joint order, FK target frames, and Jacobian columns. + +## Goals / Non-Goals + +**Goals:** +- Implement group FK/Jacobian support in Drake and RoboPlan worlds. +- Make WorldMonitor expose group-scoped state and query helpers. +- Validate group-local joint ordering and global/local joint-state mapping. +- Keep base pose/base link semantics safe by rejecting unsupported backend cases rather than silently producing wrong transforms. + +**Non-Goals:** +- Do not migrate planner or IK algorithms yet. +- Do not expose new public module APIs yet. +- Do not include Viser UI changes. + +## Decisions + +- Robot-scoped FK/Jacobian compatibility wrappers may remain temporarily but must resolve through exactly one pose-targetable group or fail clearly. +- RoboPlan Jacobians must be projected into group-local joint order; returning a full backend Jacobian is not sufficient. +- Backends should reject unsupported base-pose/base-link combinations explicitly. + +## Risks / Trade-offs + +- Backend capabilities differ. Tests should pin the supported subset and error messages. +- Some existing tests may need compatibility shims until PR 3 migrates algorithms and module APIs. diff --git a/openspec/changes/planning-world-monitor-groups/proposal.md b/openspec/changes/planning-world-monitor-groups/proposal.md new file mode 100644 index 0000000000..2d8d1681d1 --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/proposal.md @@ -0,0 +1,26 @@ +## Why + +Once planning groups exist, the backend query layer must answer group-scoped FK, Jacobian, joint-state, and collision questions. This is the second reviewable PR in the stack: it consumes the foundation without exposing the public `ManipulationModule` API or UI rewrite. + +## What Changes + +- Teach Drake and RoboPlan world backends to resolve planning groups. +- Add group FK/Jacobian APIs and group-local joint ordering behavior. +- Update world monitor and robot state monitor logic for group-aware state access. +- Preserve legacy robot-scoped wrappers only where needed for compatibility, resolving through a unique pose-targetable group. +- Add backend and monitor tests. + +## Capabilities + +### New Capabilities +- `planning-world-monitor-groups`: World and monitor support for planning-group queries. + +### Modified Capabilities +- `manipulation-world-backends`: Drake/RoboPlan backends understand group-local chains and target frames. + +## Impact + +- Base branch: PR 1 `planning-groups-foundation`. +- Reference implementation: `cc/spec/movegroup`. +- Primary files: `dimos/manipulation/planning/world/*`, `dimos/manipulation/planning/monitor/*`, RoboPlan world tests, Drake planning-group tests. +- Out of scope: IK/RRT algorithm migration, `ManipulationModule`, Viser UI, and control integration. diff --git a/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md b/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md new file mode 100644 index 0000000000..3bff6f8bd0 --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md @@ -0,0 +1,22 @@ +## ADDED Requirements + +### Requirement: World backends must answer group pose queries +World backends MUST provide end-effector pose queries for pose-targetable planning groups using the group's configured target frame. + +#### Scenario: Group pose is requested +- **WHEN** a caller requests FK for a valid pose-targetable group ID +- **THEN** the backend returns the pose of that group's `tip_link` + +### Requirement: World backends must return group-ordered Jacobians +World backends MUST return Jacobian columns ordered according to the requested group's local joint names. + +#### Scenario: Backend returns a full robot Jacobian +- **WHEN** a group contains a subset of robot joints in a custom order +- **THEN** the backend projects and reorders Jacobian columns to match the group + +### Requirement: Monitors must expose group-scoped state +World monitors MUST support planning-group state queries without requiring callers to infer robot-local joint mappings. + +#### Scenario: Group state is requested for a multi-group robot +- **WHEN** a caller requests state for a specific group ID +- **THEN** the monitor returns joint data scoped and ordered for that group diff --git a/openspec/changes/planning-world-monitor-groups/tasks.md b/openspec/changes/planning-world-monitor-groups/tasks.md new file mode 100644 index 0000000000..acb647d94a --- /dev/null +++ b/openspec/changes/planning-world-monitor-groups/tasks.md @@ -0,0 +1,18 @@ +## 1. Extract backend group support + +- [ ] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. +- [ ] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. +- [ ] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. +- [ ] 1.4 Extract world monitor and robot state monitor group-aware query changes. + +## 2. Tests + +- [ ] 2.1 Bring over Drake group world tests. +- [ ] 2.2 Bring over RoboPlan world tests, including planning split file if needed. +- [ ] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. +- [ ] 3.2 Run targeted mypy on changed backend/monitor files. +- [ ] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. diff --git a/openspec/changes/viser-group-planning-ui/.openspec.yaml b/openspec/changes/viser-group-planning-ui/.openspec.yaml new file mode 100644 index 0000000000..34f9314d22 --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-29 diff --git a/openspec/changes/viser-group-planning-ui/README.md b/openspec/changes/viser-group-planning-ui/README.md new file mode 100644 index 0000000000..72628be56b --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/README.md @@ -0,0 +1,3 @@ +# viser-group-planning-ui + +Extract PR 5 from reference branch cc/spec/movegroup: make Viser visualization and panel workflows group-aware. diff --git a/openspec/changes/viser-group-planning-ui/design.md b/openspec/changes/viser-group-planning-ui/design.md new file mode 100644 index 0000000000..0e201b7004 --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/design.md @@ -0,0 +1,28 @@ +## Context + +Viser must let users select planning groups, move joint and pose targets, see feasibility feedback, preview paths, and execute only when the preview still matches the current robot state. + +## Goals / Non-Goals + +**Goals:** +- Expose group selection and group-aware robot controls. +- Keep target ghost and preview animation tied to selected groups. +- Preserve safe execution behavior and clear recoverable errors. +- Keep UI review isolated from planner correctness review. + +**Non-Goals:** +- Do not change core planning algorithms in this PR. +- Do not alter the planning-group data model. +- Do not include control/coordinator changes. + +## Decisions + +- Use a panel backend boundary so UI code does not directly own all manipulation-module orchestration. +- Treat planning-group IDs as the UI's stable selection keys. +- Validate base-link/root assumptions before rendering URDFs under base poses. +- Keep manual demo/checklist coverage for UI feel in addition to unit tests. + +## Risks / Trade-offs + +- UI tests can pass while interaction feels wrong; require a small manual checklist before review. +- This PR is still large, but isolating it keeps visual review focused. diff --git a/openspec/changes/viser-group-planning-ui/proposal.md b/openspec/changes/viser-group-planning-ui/proposal.md new file mode 100644 index 0000000000..54137b817c --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/proposal.md @@ -0,0 +1,25 @@ +## Why + +The Viser changes are large and have different review criteria from planning algorithms. They should be reviewed as a dedicated UI/backend PR after the public manipulation group APIs exist. + +## What Changes + +- Make the Viser panel group-aware. +- Split panel/backend behavior out of the old adapter shape. +- Update scene previews, target ghosts, group selection, feasibility state, and safe execution checks. +- Update Viser tests and manual review checklist. + +## Capabilities + +### New Capabilities +- `viser-group-planning-ui`: Viser visualization supports group-aware planning, preview, target evaluation, and execution state. + +### Modified Capabilities +- `manipulation-visualization`: Visualization targets explicit planning groups rather than a robot-scoped end-effector field. + +## Impact + +- Base branch: PR 4 `manipulation-module-group-api`. +- Reference implementation: `cc/spec/movegroup`. +- Primary files: `dimos/manipulation/visualization/viser/*`, `dimos/manipulation/visualization/types.py`, `dimos/manipulation/visualization/test_factory.py`. +- Out of scope: planning model/backend/algorithm changes except as already supplied by earlier PRs. diff --git a/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md b/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md new file mode 100644 index 0000000000..a653991fcd --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md @@ -0,0 +1,22 @@ +## ADDED Requirements + +### Requirement: Viser panel must select planning groups explicitly +The Viser planning panel MUST present and use planning-group selections for group-aware joint and pose planning controls. + +#### Scenario: User selects a manipulator group +- **WHEN** a user selects a planning group in the panel +- **THEN** joint controls, pose target controls, and preview state apply to that group + +### Requirement: Viser preview must reflect group feasibility and current state +The Viser UI MUST show target feasibility and preview validity for the selected planning group. + +#### Scenario: Target is infeasible +- **WHEN** target evaluation reports infeasible for the selected group +- **THEN** the target ghost and panel state indicate that the target cannot be executed safely + +### Requirement: Viser execution must require a fresh matching plan +The Viser panel MUST prevent execution when no fresh plan matches the current selected robot/group state. + +#### Scenario: Current state no longer matches preview +- **WHEN** the robot state changes after planning +- **THEN** the execute action is rejected until a fresh matching plan is generated diff --git a/openspec/changes/viser-group-planning-ui/tasks.md b/openspec/changes/viser-group-planning-ui/tasks.md new file mode 100644 index 0000000000..0e9523358e --- /dev/null +++ b/openspec/changes/viser-group-planning-ui/tasks.md @@ -0,0 +1,17 @@ +## 1. Extract Viser UI changes + +- [ ] 1.1 Start from PR 4 and use `cc/spec/movegroup` as reference. +- [ ] 1.2 Extract Viser panel/backend/scene/state/visualizer group-aware changes. +- [ ] 1.3 Remove or replace the old adapter code only as required by this UI slice. +- [ ] 1.4 Extract visualization type/factory updates needed by Viser. + +## 2. Tests and manual check + +- [ ] 2.1 Bring over Viser unit tests and lifecycle tests. +- [ ] 2.2 Run the manual checklist: group selector, joint sliders, pose gizmo, target ghost, infeasible target color, path preview, execute gate, clear path. + +## 3. Validation + +- [ ] 3.1 Run `uv run pytest dimos/manipulation/visualization/test_factory.py dimos/manipulation/visualization/viser/test_*.py -q`. +- [ ] 3.2 Run targeted mypy on changed Viser production files if practical. +- [ ] 3.3 Verify no planner/backend implementation changes are included beyond compile fixes against PR 4. diff --git a/openspec/config.yaml b/openspec/config.yaml new file mode 100644 index 0000000000..392946c67c --- /dev/null +++ b/openspec/config.yaml @@ -0,0 +1,20 @@ +schema: spec-driven + +# Project context (optional) +# This is shown to AI when creating artifacts. +# Add your tech stack, conventions, style guides, domain knowledge, etc. +# Example: +# context: | +# Tech stack: TypeScript, React, Node.js +# We use conventional commits +# Domain: e-commerce platform + +# Per-artifact rules (optional) +# Add custom rules for specific artifacts. +# Example: +# rules: +# proposal: +# - Keep proposals under 500 words +# - Always include a "Non-goals" section +# tasks: +# - Break tasks into chunks of max 2 hours From 842f70a9971545a5881dfd7cda25f1cb9154a850 Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:59:49 -0700 Subject: [PATCH 02/29] feat: Add planning groups foundation (#2644) --- .../manipulation/planning/groups/discovery.py | 352 +++++++++++ .../planning/groups/identifiers.py | 112 ++++ dimos/manipulation/planning/groups/models.py | 112 ++++ .../manipulation/planning/groups/registry.py | 108 ++++ .../planning/groups/test_planning_groups.py | 571 ++++++++++++++++++ dimos/manipulation/planning/groups/utils.py | 143 +++++ .../planning/kinematics/test_pink_ik.py | 10 +- .../planning/monitor/test_world_monitor.py | 7 +- dimos/manipulation/planning/spec/config.py | 67 +- dimos/manipulation/planning/spec/models.py | 9 + .../manipulation/test_manipulation_module.py | 10 +- dimos/manipulation/test_manipulation_unit.py | 37 +- dimos/manipulation/test_planning_factory.py | 10 +- dimos/manipulation/test_roboplan_world.py | 10 +- .../visualization/test_factory.py | 7 +- .../viser/test_visualizer_lifecycle.py | 7 +- dimos/robot/manipulators/_modeling.py | 11 +- dimos/robot/manipulators/a750/config.py | 13 +- dimos/robot/manipulators/openarm/config.py | 25 +- dimos/robot/manipulators/piper/config.py | 13 +- dimos/robot/manipulators/xarm/config.py | 20 +- .../manipulation/adding_a_custom_arm.md | 37 +- .../manipulation/planning_groups.md | 161 +++++ .../planning-groups-foundation/proposal.md | 4 +- .../planning-groups-foundation/tasks.md | 20 +- 25 files changed, 1811 insertions(+), 65 deletions(-) create mode 100644 dimos/manipulation/planning/groups/discovery.py create mode 100644 dimos/manipulation/planning/groups/identifiers.py create mode 100644 dimos/manipulation/planning/groups/models.py create mode 100644 dimos/manipulation/planning/groups/registry.py create mode 100644 dimos/manipulation/planning/groups/test_planning_groups.py create mode 100644 dimos/manipulation/planning/groups/utils.py create mode 100644 docs/capabilities/manipulation/planning_groups.md diff --git a/dimos/manipulation/planning/groups/discovery.py b/dimos/manipulation/planning/groups/discovery.py new file mode 100644 index 0000000000..49cad388ff --- /dev/null +++ b/dimos/manipulation/planning/groups/discovery.py @@ -0,0 +1,352 @@ +# 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. + +"""Planning group discovery from SRDF or conservative model fallback.""" + +from __future__ import annotations + +import itertools +from pathlib import Path +import xml.etree.ElementTree as ET + +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.robot.model_parser import JointDescription, ModelDescription +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +FALLBACK_PLANNING_GROUP_NAME = "manipulator" + + +class PlanningGroupDiscoveryError(ValueError): + """Raised when planning groups cannot be discovered for a model.""" + + +def discover_planning_group_definitions( + *, + robot_name: str, + model_path: Path, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path | None = None, +) -> list[PlanningGroupDefinition]: + """Discover planning groups from SRDF or fallback generation. + + Precedence is explicit SRDF path, conservative auto-discovery with warning, + then fallback generation from the controllable joint set. + """ + resolved_srdf_path = _resolve_srdf_path(model_path, srdf_path) + if resolved_srdf_path is not None: + groups = parse_srdf_planning_groups( + resolved_srdf_path, + model=model, + controllable_joint_names=controllable_joint_names, + ) + if groups: + return groups + logger.warning( + f"No supported planning groups found in SRDF {resolved_srdf_path} " + f"for robot {robot_name}; trying fallback generation" + ) + + return [ + generate_fallback_planning_group( + model=model, + controllable_joint_names=controllable_joint_names, + ) + ] + + +def parse_srdf_planning_groups( + srdf_path: Path, + *, + model: ModelDescription, + controllable_joint_names: list[str], +) -> list[PlanningGroupDefinition]: + """Extract supported SRDF planning group definitions. + + Supported forms are a single ```` + child or an ordered list of ```` children. Other forms, + including SRDF ```` metadata, are ignored for planning group + extraction. This is intentionally a minimal SRDF group extractor rather + than a full SRDF parser; adopting a ROS/MoveIt parser such as srdfdom would + add substantial dependency overhead for this narrow subset. + """ + root = ET.parse(srdf_path).getroot() + groups: list[PlanningGroupDefinition] = [] + for group_elem in root.findall("group"): + group_name = group_elem.get("name") + if not group_name: + logger.warning(f"Skipping SRDF group without a name in {srdf_path}") + continue + + children = [child for child in list(group_elem) if isinstance(child.tag, str)] + chain_children = [child for child in children if child.tag == "chain"] + joint_children = [child for child in children if child.tag == "joint"] + unsupported_children = [child for child in children if child.tag not in {"chain", "joint"}] + + if len(chain_children) == 1 and not joint_children and not unsupported_children: + definition = _parse_chain_group( + group_name, + chain_children[0], + model=model, + controllable_joint_names=controllable_joint_names, + srdf_path=srdf_path, + ) + elif joint_children and len(joint_children) == len(children): + definition = _parse_joint_list_group( + group_name, + joint_children, + model=model, + controllable_joint_names=controllable_joint_names, + srdf_path=srdf_path, + ) + else: + child_tags = [child.tag for child in children] + logger.warning( + f"Skipping unsupported SRDF planning group {group_name} in " + f"{srdf_path} with child tags {child_tags}" + ) + definition = None + + if definition is not None: + groups.append(definition) + + return groups + + +def generate_fallback_planning_group( + *, + model: ModelDescription, + controllable_joint_names: list[str], +) -> PlanningGroupDefinition: + """Generate one conservative fallback planning group named ``manipulator``.""" + ordered_joints = _validate_and_order_serial_joints(model, controllable_joint_names) + while ordered_joints and ordered_joints[-1].type == "prismatic": + removed = ordered_joints.pop() + logger.warning( + f"Excluding terminal prismatic joint {removed.name} from " + f"fallback planning group {FALLBACK_PLANNING_GROUP_NAME}" + ) + + if not ordered_joints: + raise PlanningGroupDiscoveryError( + "Fallback planning group generation removed all candidate joints; provide SRDF" + ) + + return PlanningGroupDefinition( + name=FALLBACK_PLANNING_GROUP_NAME, + joint_names=tuple(joint.name for joint in ordered_joints), + base_link=ordered_joints[0].parent_link, + tip_link=ordered_joints[-1].child_link, + source="fallback", + ) + + +def _resolve_srdf_path(model_path: Path, srdf_path: Path | None) -> Path | None: + if srdf_path is not None: + if srdf_path.exists(): + return srdf_path + raise FileNotFoundError(f"SRDF file not found: {srdf_path}") + + for candidate in _srdf_auto_discovery_candidates(model_path): + if candidate.exists(): + logger.warning(f"Auto-discovered SRDF at {candidate}") + return candidate + return None + + +def _srdf_auto_discovery_candidates(model_path: Path) -> list[Path]: + candidates: list[Path] = [] + name = model_path.name + if name.endswith(".urdf.xacro"): + candidates.append(model_path.with_name(name.removesuffix(".urdf.xacro") + ".srdf")) + elif model_path.suffix: + candidates.append(model_path.with_suffix(".srdf")) + candidates.append(model_path.parent / "config" / "robot.srdf") + candidates.append(model_path.parent.parent / "config" / "robot.srdf") + return list(dict.fromkeys(candidates)) + + +def _parse_chain_group( + group_name: str, + chain_elem: ET.Element, + *, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path, +) -> PlanningGroupDefinition | None: + base_link = chain_elem.get("base_link") + tip_link = chain_elem.get("tip_link") + if not base_link or not tip_link: + logger.warning( + f"Skipping SRDF chain group {group_name} in {srdf_path} because " + "base_link or tip_link is missing" + ) + return None + + try: + ordered_joints = _ordered_joints_between_links(model, base_link, tip_link) + controlled_joints = [joint for joint in ordered_joints if joint.type != "fixed"] + _validate_controllable(group_name, controlled_joints, controllable_joint_names) + except PlanningGroupDiscoveryError as exc: + logger.warning(f"Skipping SRDF chain group {group_name} in {srdf_path}: {exc}") + return None + + return PlanningGroupDefinition( + name=group_name, + joint_names=tuple(joint.name for joint in controlled_joints), + base_link=base_link, + tip_link=tip_link, + ) + + +def _parse_joint_list_group( + group_name: str, + joint_children: list[ET.Element], + *, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path, +) -> PlanningGroupDefinition | None: + joint_names = [child.get("name", "") for child in joint_children] + if any(not name for name in joint_names): + logger.warning( + f"Skipping SRDF joint-list group {group_name} in {srdf_path} with empty joint name" + ) + return None + try: + ordered_joints = _validate_ordered_serial_joints(model, joint_names) + _validate_controllable(group_name, ordered_joints, controllable_joint_names) + except PlanningGroupDiscoveryError as exc: + logger.warning(f"Skipping SRDF joint-list group {group_name} in {srdf_path}: {exc}") + return None + + return PlanningGroupDefinition( + name=group_name, + joint_names=tuple(joint.name for joint in ordered_joints), + base_link=ordered_joints[0].parent_link, + tip_link=ordered_joints[-1].child_link, + ) + + +def _ordered_joints_between_links( + model: ModelDescription, + base_link: str, + tip_link: str, +) -> list[JointDescription]: + joints_by_parent: dict[str, list[JointDescription]] = {} + for joint in model.joints: + joints_by_parent.setdefault(joint.parent_link, []).append(joint) + + ordered_joints: list[JointDescription] = [] + current_link = base_link + visited_links = {base_link} + while current_link != tip_link: + children = joints_by_parent.get(current_link, []) + if len(children) != 1: + raise PlanningGroupDiscoveryError( + f"chain from {base_link} to {tip_link} is branching or disconnected at {current_link}" + ) + joint = children[0] + ordered_joints.append(joint) + current_link = joint.child_link + if current_link in visited_links: + raise PlanningGroupDiscoveryError("chain contains a cycle") + visited_links.add(current_link) + + return ordered_joints + + +def _validate_ordered_serial_joints( + model: ModelDescription, + joint_names: list[str], +) -> list[JointDescription]: + ordered_joints: list[JointDescription] = [] + for joint_name in joint_names: + joint = model.get_joint(joint_name) + if joint is None: + raise PlanningGroupDiscoveryError(f"joint {joint_name} does not exist in model") + if joint.type == "fixed": + raise PlanningGroupDiscoveryError(f"joint {joint_name} is fixed") + ordered_joints.append(joint) + + if not ordered_joints: + raise PlanningGroupDiscoveryError("planning group contains no joints") + + for previous, current in itertools.pairwise(ordered_joints): + if previous.child_link != current.parent_link: + raise PlanningGroupDiscoveryError( + f"joints {previous.name} and {current.name} are not adjacent in a serial chain" + ) + return ordered_joints + + +def _validate_and_order_serial_joints( + model: ModelDescription, + joint_names: list[str], +) -> list[JointDescription]: + if not joint_names: + raise PlanningGroupDiscoveryError("fallback requires at least one controllable joint") + + joints: list[JointDescription] = [] + for joint_name in joint_names: + joint = model.get_joint(joint_name) + if joint is None: + raise PlanningGroupDiscoveryError(f"joint {joint_name} does not exist in model") + if joint.type == "fixed": + raise PlanningGroupDiscoveryError(f"joint {joint_name} is fixed") + joints.append(joint) + + by_parent = {joint.parent_link: joint for joint in joints} + by_child = {joint.child_link: joint for joint in joints} + if len(by_parent) != len(joints) or len(by_child) != len(joints): + raise PlanningGroupDiscoveryError("controllable joints branch or merge; provide SRDF") + + starts = [joint for joint in joints if joint.parent_link not in by_child] + ends = [joint for joint in joints if joint.child_link not in by_parent] + if len(starts) != 1 or len(ends) != 1: + raise PlanningGroupDiscoveryError( + "controllable joints are disconnected or cyclic; provide SRDF" + ) + + ordered_joints: list[JointDescription] = [] + current = starts[0] + while True: + ordered_joints.append(current) + next_joint = by_parent.get(current.child_link) + if next_joint is None: + break + current = next_joint + + if len(ordered_joints) != len(joints): + raise PlanningGroupDiscoveryError("controllable joints are disconnected; provide SRDF") + return ordered_joints + + +def _validate_controllable( + group_name: str, + joints: list[JointDescription], + controllable_joint_names: list[str], +) -> None: + if not joints: + raise PlanningGroupDiscoveryError( + f"planning group {group_name} contains no controllable joints" + ) + controllable = set(controllable_joint_names) + missing = [joint.name for joint in joints if joint.name not in controllable] + if missing: + raise PlanningGroupDiscoveryError( + f"planning group {group_name} includes joints outside controllable set: {missing}" + ) diff --git a/dimos/manipulation/planning/groups/identifiers.py b/dimos/manipulation/planning/groups/identifiers.py new file mode 100644 index 0000000000..db9d990081 --- /dev/null +++ b/dimos/manipulation/planning/groups/identifiers.py @@ -0,0 +1,112 @@ +# 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. + +"""Planning-group and global-joint identifier helpers.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + LocalModelJointName, + PlanningGroupID, + RobotName, +) + + +def assert_valid_robot_name(robot_name: RobotName) -> None: + """Validate a robot name for delimiter-based public IDs.""" + if not robot_name or "/" in robot_name: + raise ValueError(f"Invalid robot name: {robot_name!r}") + + +def assert_valid_local_joint_name(local_joint_name: LocalModelJointName) -> None: + """Validate a local model joint name for delimiter-based global joint names.""" + if not local_joint_name or "/" in local_joint_name: + raise ValueError(f"Invalid local joint name: {local_joint_name!r}") + + +def assert_local_joint_names(names: Sequence[LocalModelJointName]) -> None: + """Validate that names are local model joint names, not global joint names.""" + for name in names: + assert_valid_local_joint_name(name) + + +def make_planning_group_id(robot_name: RobotName, group_name: str) -> PlanningGroupID: + """Build a public planning group ID.""" + assert_valid_robot_name(robot_name) + if not group_name or "/" in group_name: + raise ValueError(f"Invalid planning group name: {group_name!r}") + return f"{robot_name}/{group_name}" + + +def parse_planning_group_id(group_id: PlanningGroupID) -> tuple[RobotName, str]: + """Split and validate a planning group ID.""" + parts = group_id.split("/", maxsplit=1) + if len(parts) != 2 or not parts[0] or not parts[1] or "/" in parts[1]: + raise ValueError( + f"Invalid planning group ID {group_id!r}; expected '{{robot_name}}/{{group_name}}'" + ) + return parts[0], parts[1] + + +def make_global_joint_name( + robot_name: RobotName, + local_joint_name: LocalModelJointName, +) -> GlobalJointName: + """Convert a local model joint name to a public global joint name.""" + assert_valid_robot_name(robot_name) + assert_valid_local_joint_name(local_joint_name) + return f"{robot_name}/{local_joint_name}" + + +def make_global_joint_names( + robot_name: RobotName, + local_joint_names: list[LocalModelJointName] | tuple[LocalModelJointName, ...], +) -> list[GlobalJointName]: + """Convert local model joint names to public global joint names.""" + return [make_global_joint_name(robot_name, name) for name in local_joint_names] + + +def is_global_joint_name(name: str) -> bool: + """Return whether name has the exact global joint-name shape.""" + parts = name.split("/") + return len(parts) == 2 and bool(parts[0]) and bool(parts[1]) + + +def assert_global_joint_names(names: Sequence[GlobalJointName]) -> None: + """Validate that names are global joint names.""" + invalid = [name for name in names if not is_global_joint_name(name)] + if invalid: + raise ValueError(f"Expected global joint names; got invalid names: {invalid}") + + +def local_joint_name_from_global( + robot_name: RobotName, + global_joint_name: GlobalJointName, +) -> LocalModelJointName: + """Validate and strip a global joint name for backend internals.""" + assert_valid_robot_name(robot_name) + prefix = f"{robot_name}/" + if not global_joint_name.startswith(prefix): + raise ValueError( + f"Global joint name {global_joint_name!r} does not belong to robot {robot_name!r}" + ) + local_name = global_joint_name[len(prefix) :] + try: + assert_valid_local_joint_name(local_name) + except ValueError as exc: + raise ValueError(f"Invalid global joint name: {global_joint_name!r}") from exc + return local_name diff --git a/dimos/manipulation/planning/groups/models.py b/dimos/manipulation/planning/groups/models.py new file mode 100644 index 0000000000..c08b1bd4b5 --- /dev/null +++ b/dimos/manipulation/planning/groups/models.py @@ -0,0 +1,112 @@ +# 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. + +"""Backend-independent planning-group domain models.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, TypeAlias + +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + LocalModelJointName, + PlanningGroupID, + RobotName, +) + +PlanningGroupSource: TypeAlias = Literal["srdf", "fallback"] + + +@dataclass(frozen=True) +class PlanningGroupDefinition: + """Model-level declaration of a planning group. + + Joint names are local model names. The definition is safe to store on + ``RobotModelConfig`` and is not bound to any runtime world robot ID. + """ + + name: str + joint_names: tuple[LocalModelJointName, ...] + base_link: str + tip_link: str | None = None + source: PlanningGroupSource = "srdf" + + @property + def has_pose_target(self) -> bool: + """Whether this group has a valid pose target frame.""" + return self.tip_link is not None + + +@dataclass(frozen=True) +class PlanningGroup: + """Public backend-independent planning group. + + A planning group exposes stable public IDs and global joint names for + planning APIs. It intentionally does not include backend runtime robot IDs. + """ + + id: PlanningGroupID + robot_name: RobotName + group_name: str + joint_names: tuple[GlobalJointName, ...] + local_joint_names: tuple[LocalModelJointName, ...] + base_link: str + tip_link: str | None = None + source: PlanningGroupSource = "srdf" + + @property + def has_pose_target(self) -> bool: + """Whether this group can be directly pose-targeted.""" + return self.tip_link is not None + + +@dataclass(frozen=True) +class PlanningGroupSelection: + """Validated ordered selection of planning groups. + + Selection validates ID existence and selected-joint overlap outside any + world backend. Requested group order is preserved. + """ + + groups: tuple[PlanningGroup, ...] + group_ids: tuple[PlanningGroupID, ...] + joint_names: tuple[GlobalJointName, ...] + robot_names: tuple[RobotName, ...] + + @classmethod + def from_groups(cls, groups: tuple[PlanningGroup, ...]) -> PlanningGroupSelection: + """Build a selection, rejecting overlapping selected global joints.""" + seen_joints: dict[GlobalJointName, PlanningGroupID] = {} + joint_names: list[GlobalJointName] = [] + robot_names: list[RobotName] = [] + for group in groups: + if group.robot_name not in robot_names: + robot_names.append(group.robot_name) + for joint_name in group.joint_names: + previous_group_id = seen_joints.get(joint_name) + if previous_group_id is not None: + raise ValueError( + "Selected planning groups overlap on global joint " + f"{joint_name}: {previous_group_id} and {group.id}" + ) + seen_joints[joint_name] = group.id + joint_names.append(joint_name) + + return cls( + groups=groups, + group_ids=tuple(group.id for group in groups), + joint_names=tuple(joint_names), + robot_names=tuple(robot_names), + ) diff --git a/dimos/manipulation/planning/groups/registry.py b/dimos/manipulation/planning/groups/registry.py new file mode 100644 index 0000000000..c5887a090e --- /dev/null +++ b/dimos/manipulation/planning/groups/registry.py @@ -0,0 +1,108 @@ +# 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. + +"""Backend-independent planning-group registry.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +from dimos.manipulation.planning.groups.discovery import FALLBACK_PLANNING_GROUP_NAME +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, +) +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.spec.models import PlanningGroupID, RobotName + +if TYPE_CHECKING: + from dimos.manipulation.planning.spec.config import RobotModelConfig + + +class PlanningGroupRegistry: + """Registry of public planning groups derived from robot configs.""" + + def __init__(self, robot_configs: Iterable[RobotModelConfig] = ()) -> None: + self._groups: dict[PlanningGroupID, PlanningGroup] = {} + self._groups_by_robot: dict[RobotName, list[PlanningGroup]] = {} + for config in robot_configs: + self.add_robot(config) + + def add_robot(self, config: RobotModelConfig) -> None: + """Register all planning groups declared by one robot config.""" + if config.name in self._groups_by_robot: + raise ValueError(f"Robot '{config.name}' is already registered") + + robot_groups: list[PlanningGroup] = [] + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if group_id in self._groups: + raise ValueError(f"Planning group '{group_id}' is already registered") + group = PlanningGroup( + id=group_id, + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + self._groups[group_id] = group + robot_groups.append(group) + self._groups_by_robot[config.name] = robot_groups + + def list(self) -> tuple[PlanningGroup, ...]: + """List planning groups in robot registration order.""" + groups: list[PlanningGroup] = [] + for robot_groups in self._groups_by_robot.values(): + groups.extend(robot_groups) + return tuple(groups) + + def get(self, group_id: PlanningGroupID) -> PlanningGroup: + """Return one planning group by public ID.""" + try: + return self._groups[group_id] + except KeyError as exc: + raise KeyError(f"Unknown planning group ID: {group_id}") from exc + + def select(self, group_ids: Iterable[PlanningGroupID]) -> PlanningGroupSelection: + """Validate and return an ordered planning-group selection.""" + return PlanningGroupSelection.from_groups( + tuple(self.get(group_id) for group_id in group_ids) + ) + + def groups_for_robot(self, robot_name: RobotName) -> tuple[PlanningGroup, ...]: + """Return planning groups for one robot.""" + return tuple(self._groups_by_robot.get(robot_name, ())) + + def default_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: + """Return the generated fallback group ID for robot-scoped wrappers.""" + group_id = make_planning_group_id(robot_name, FALLBACK_PLANNING_GROUP_NAME) + return group_id if group_id in self._groups else None + + def primary_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: + """Return the unique pose-targetable group ID for robot-scoped wrappers.""" + pose_groups = [ + group for group in self.groups_for_robot(robot_name) if group.has_pose_target + ] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError( + f"Robot '{robot_name}' has {len(pose_groups)} pose-targetable planning groups; " + "use an explicit planning group ID" + ) + return pose_groups[0].id diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py new file mode 100644 index 0000000000..ae5776fd74 --- /dev/null +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -0,0 +1,571 @@ +# 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. + +"""Tests for planning groups.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from dimos.manipulation.planning.groups import utils as planning_group_utils +from dimos.manipulation.planning.groups.discovery import ( + FALLBACK_PLANNING_GROUP_NAME, + PlanningGroupDiscoveryError, + discover_planning_group_definitions, + generate_fallback_planning_group, + parse_srdf_planning_groups, +) +from dimos.manipulation.planning.groups.identifiers import local_joint_name_from_global +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + filter_joint_state_to_selected_joints, + joint_target_to_global_names, + matching_global_joint_name, + planning_group_id_from_selector, +) +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.model_parser import JointDescription, ModelDescription + + +def _serial_model(*joint_types: str) -> ModelDescription: + joints = [ + JointDescription( + name=f"joint{i + 1}", + type=joint_type, + parent_link=f"link{i}", + child_link=f"link{i + 1}", + ) + for i, joint_type in enumerate(joint_types) + ] + return ModelDescription( + joints=joints, + root_link="link0", + links=[f"link{i}" for i in range(len(joint_types) + 1)], + ) + + +def _branching_model() -> ModelDescription: + return ModelDescription( + joints=[ + JointDescription( + name="left_joint", + type="revolute", + parent_link="base", + child_link="left_link", + ), + JointDescription( + name="right_joint", + type="revolute", + parent_link="base", + child_link="right_link", + ), + ], + root_link="base", + links=["base", "left_link", "right_link"], + ) + + +def _write_srdf(tmp_path: Path, body: str) -> Path: + srdf_path = tmp_path / "robot.srdf" + srdf_path.write_text(f"{body}") + return srdf_path + + +def _make_group() -> PlanningGroup: + return PlanningGroup( + id="left/arm", + robot_name="left", + group_name="arm", + joint_names=("left/j1", "left/j2", "left/j3"), + local_joint_names=("j1", "j2", "j3"), + base_link="base", + tip_link="ee", + ) + + +def _robot_config( + name: str = "robot", + planning_groups: list[PlanningGroupDefinition] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/tmp/robot.urdf"), + base_pose=PoseStamped(), + joint_names=["joint1", "joint2", "joint3"], + planning_groups=planning_groups + if planning_groups is not None + else [ + PlanningGroupDefinition( + name=FALLBACK_PLANNING_GROUP_NAME, + joint_names=("joint1", "joint2"), + base_link="base", + tip_link="tool", + ) + ], + ) + + +def test_parse_srdf_chain_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + "", + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert len(groups) == 1 + assert groups[0].name == "arm" + assert groups[0].joint_names == ("joint1", "joint2", "joint3") + assert groups[0].base_link == "link0" + assert groups[0].tip_link == "link3" + assert groups[0].source == "srdf" + + +def test_parse_srdf_ordered_joint_list_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "prismatic", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert len(groups) == 1 + assert groups[0].joint_names == ("joint1", "joint2", "joint3") + assert groups[0].base_link == "link0" + assert groups[0].tip_link == "link3" + + +def test_parse_srdf_skips_unsupported_groups_and_ignores_end_effector( + tmp_path: Path, +) -> None: + model = _serial_model("revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + ) + + assert [group.name for group in groups] == ["arm"] + + +def test_fallback_generates_manipulator_for_unambiguous_serial_chain() -> None: + model = _serial_model("revolute", "prismatic", "revolute") + + group = generate_fallback_planning_group( + model=model, + controllable_joint_names=["joint2", "joint1", "joint3"], + ) + + assert group.name == FALLBACK_PLANNING_GROUP_NAME + assert group.joint_names == ("joint1", "joint2", "joint3") + assert group.base_link == "link0" + assert group.tip_link == "link3" + assert group.source == "fallback" + + +def test_fallback_strips_terminal_prismatic_joints() -> None: + model = _serial_model("revolute", "revolute", "prismatic") + + group = generate_fallback_planning_group( + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert group.joint_names == ("joint1", "joint2") + assert group.tip_link == "link2" + assert group.source == "fallback" + + +def test_fallback_rejects_branching_model() -> None: + with pytest.raises(PlanningGroupDiscoveryError, match="branch"): + generate_fallback_planning_group( + model=_branching_model(), + controllable_joint_names=["left_joint", "right_joint"], + ) + + +def test_fallback_rejects_all_terminal_prismatic_candidates() -> None: + with pytest.raises(PlanningGroupDiscoveryError, match="removed all candidate joints"): + generate_fallback_planning_group( + model=_serial_model("prismatic", "prismatic"), + controllable_joint_names=["joint1", "joint2"], + ) + + +def test_parse_srdf_skips_invalid_groups_and_keeps_valid_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + ) + + assert [group.name for group in groups] == ["arm"] + + +def test_discovery_rejects_missing_explicit_srdf(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="SRDF file not found"): + discover_planning_group_definitions( + robot_name="robot", + model_path=tmp_path / "robot.urdf", + model=_serial_model("revolute"), + controllable_joint_names=["joint1"], + srdf_path=tmp_path / "missing.srdf", + ) + + +def test_discovery_falls_back_when_srdf_has_no_supported_groups(tmp_path: Path) -> None: + model_path = tmp_path / "robot.urdf.xacro" + model_path.write_text("") + (tmp_path / "robot.srdf").write_text( + "" + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=_serial_model("revolute"), + controllable_joint_names=["joint1"], + ) + + assert [group.name for group in groups] == [FALLBACK_PLANNING_GROUP_NAME] + assert [group.source for group in groups] == ["fallback"] + + +def test_discovery_prefers_explicit_srdf_over_fallback(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute") + model_path = tmp_path / "robot.urdf" + model_path.write_text("") + srdf_path = _write_srdf( + tmp_path, + "", + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + srdf_path=srdf_path, + ) + + assert [group.name for group in groups] == ["srdf_arm"] + + +def test_discovery_auto_discovers_srdf(tmp_path: Path) -> None: + model = _serial_model("revolute") + model_path = tmp_path / "robot.urdf" + model_path.write_text("") + _write_srdf( + tmp_path, + "", + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=model, + controllable_joint_names=["joint1"], + ) + + assert [group.name for group in groups] == ["auto_arm"] + + +def test_primary_pose_group_id_for_robot_raises_when_ambiguous() -> None: + registry = PlanningGroupRegistry( + [ + RobotModelConfig( + name="robot", + model_path=Path("/tmp/robot.urdf"), + base_pose=PoseStamped(), + joint_names=["joint1", "joint2"], + planning_groups=[ + PlanningGroupDefinition( + name="left", + joint_names=("joint1",), + base_link="base", + tip_link="left_tool", + ), + PlanningGroupDefinition( + name="right", + joint_names=("joint2",), + base_link="base", + tip_link="right_tool", + ), + ], + ) + ] + ) + + with pytest.raises(ValueError, match="multiple|2 pose-targetable|explicit planning group"): + registry.primary_pose_group_id_for_robot("robot") + + +def test_registry_preserves_order_and_exposes_defaults() -> None: + registry = PlanningGroupRegistry([_robot_config("left"), _robot_config("right")]) + + assert [group.id for group in registry.list()] == ["left/manipulator", "right/manipulator"] + assert registry.default_group_id_for_robot("left") == "left/manipulator" + assert registry.primary_pose_group_id_for_robot("right") == "right/manipulator" + assert registry.get("left/manipulator").source == "srdf" + assert registry.groups_for_robot("missing") == () + assert registry.default_group_id_for_robot("missing") is None + + +def test_registry_rejects_duplicate_robot_and_unknown_group() -> None: + registry = PlanningGroupRegistry([_robot_config()]) + + with pytest.raises(ValueError, match="already registered"): + registry.add_robot(_robot_config()) + with pytest.raises(KeyError, match="Unknown planning group ID"): + registry.get("robot/missing") + + +def test_selection_preserves_group_order_and_rejects_overlapping_joints() -> None: + registry = PlanningGroupRegistry( + [ + _robot_config( + planning_groups=[ + PlanningGroupDefinition("arm", ("joint1", "joint2"), "base", "tool"), + PlanningGroupDefinition("gripper", ("joint3",), "tool"), + ] + ) + ] + ) + + selection = registry.select(["robot/gripper", "robot/arm"]) + + assert selection.group_ids == ("robot/gripper", "robot/arm") + assert selection.joint_names == ("robot/joint3", "robot/joint1", "robot/joint2") + assert selection.robot_names == ("robot",) + + overlapping = ( + PlanningGroup("robot/first", "robot", "first", ("robot/joint1",), ("joint1",), "base"), + PlanningGroup("robot/second", "robot", "second", ("robot/joint1",), ("joint1",), "base"), + ) + with pytest.raises(ValueError, match="overlap"): + type(selection).from_groups(overlapping) + + +def test_joint_target_to_global_names_accepts_named_global_targets_in_group_order() -> None: + group = _make_group() + target = JointState({"name": ["left/j3", "left/j1", "left/j2"], "position": [3.0, 1.0, 2.0]}) + + normalized = joint_target_to_global_names(group, target) + + assert normalized.name == ["left/j1", "left/j2", "left/j3"] + assert normalized.position == [1.0, 2.0, 3.0] + + +def test_joint_target_to_global_names_accepts_named_local_targets_in_group_order() -> None: + group = _make_group() + target = JointState({"name": ["j2", "j3", "j1"], "position": [2.0, 3.0, 1.0]}) + + normalized = joint_target_to_global_names(group, target) + + assert normalized.name == ["left/j1", "left/j2", "left/j3"] + assert normalized.position == [1.0, 2.0, 3.0] + + +def test_joint_target_to_global_names_rejects_mixed_global_and_local_target_names() -> None: + group = _make_group() + target = JointState({"name": ["left/j1", "j2", "left/j3"], "position": [1.0, 2.0, 3.0]}) + + with pytest.raises(ValueError, match="mixes global and local joint names"): + joint_target_to_global_names(group, target) + + +def test_joint_target_to_global_names_rejects_bad_counts_missing_and_extra() -> None: + group = _make_group() + + with pytest.raises(ValueError, match="2 positions, expected 3"): + joint_target_to_global_names(group, JointState({"position": [1.0, 2.0]})) + with pytest.raises(ValueError, match="2 names but 3 positions"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2"], "position": [1.0, 2.0, 3.0]}) + ) + with pytest.raises(ValueError, match="missing joints"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2"], "position": [1.0, 2.0]}) + ) + with pytest.raises(ValueError, match="extra joints"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2", "j3", "j4"], "position": [1.0, 2.0, 3.0, 4.0]}) + ) + + +def test_filter_joint_state_to_selected_joints_uses_local_fallbacks() -> None: + joint_state = JointState({"name": ["j1", "robot/j2"], "position": [1.0, 2.0]}) + + filtered = filter_joint_state_to_selected_joints( + joint_state, + ["robot/j1", "robot/j2"], + ["j1", "j2"], + ) + + assert filtered.name == ["robot/j1", "robot/j2"] + assert filtered.position == [1.0, 2.0] + + +def test_filter_joint_state_to_selected_joints_rejects_mismatched_and_missing_names() -> None: + joint_state = JointState({"name": ["robot/j1"], "position": [1.0]}) + + with pytest.raises(ValueError, match="same length"): + filter_joint_state_to_selected_joints(joint_state, ["robot/j1", "robot/j2"], ["j1"]) + with pytest.raises(ValueError, match="missing selected joints"): + filter_joint_state_to_selected_joints(joint_state, ["robot/j1", "robot/j2"]) + + +def test_matching_global_joint_name_requires_unique_suffix_match() -> None: + assert matching_global_joint_name({"left/j1": 1.0, "right/j2": 2.0}, "j1") == "left/j1" + assert matching_global_joint_name({"left/j1": 1.0, "right/j1": 2.0}, "j1") is None + assert matching_global_joint_name({"left/j1": 1.0}, "j2") is None + + +def test_matching_global_joint_name_warns_when_suffix_match_is_not_unique( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: list[str] = [] + + def capture_warning(message: str) -> None: + messages.append(message) + + monkeypatch.setattr(planning_group_utils.logger, "warning", capture_warning) + + assert matching_global_joint_name({"left/j1": 1.0, "right/j1": 2.0}, "j1") is None + assert messages == [ + "Expected exactly one global joint ending with local joint name 'j1', " + "found 2 matches: ['left/j1', 'right/j1']" + ] + + +def test_planning_group_id_from_selector_accepts_id_or_group() -> None: + group = _make_group() + + assert planning_group_id_from_selector(group) == "left/arm" + assert planning_group_id_from_selector("left/arm") == "left/arm" + + +def test_local_joint_name_from_global_validates_robot_prefix_and_local_shape() -> None: + assert local_joint_name_from_global("robot", "robot/j1") == "j1" + with pytest.raises(ValueError, match="does not belong"): + local_joint_name_from_global("robot", "other/j1") + with pytest.raises(ValueError, match="Invalid global joint name"): + local_joint_name_from_global("robot", "robot/") + + +def test_robot_model_config_derives_legacy_end_effector_link_from_pose_group() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1", "j2"], + joint_name_mapping={"hw_j1": "j1", "hw_j2": "j2"}, + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j1", "j2"), + base_link="base", + tip_link="tool", + ) + ], + ) + + assert config.end_effector_link == "tool" + assert config.get_urdf_joint_name("hw_j1") == "j1" + assert config.get_coordinator_joint_name("j2") == "hw_j2" + assert config.get_coordinator_joint_names() == ["hw_j1", "hw_j2"] + + +def test_robot_model_config_end_effector_link_requires_pose_group() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1"], + planning_groups=[ + PlanningGroupDefinition( + name="joint_only", + joint_names=("j1",), + base_link="base", + ) + ], + ) + + with pytest.raises(ValueError, match="no pose-target planning group"): + _ = config.end_effector_link + + +def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1", "j2"], + planning_groups=[ + PlanningGroupDefinition( + name="left", + joint_names=("j1",), + base_link="base", + tip_link="left_tool", + ), + PlanningGroupDefinition( + name="right", + joint_names=("j2",), + base_link="base", + tip_link="right_tool", + ), + ], + ) + + with pytest.raises(ValueError, match="multiple pose-target planning groups"): + _ = config.end_effector_link diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py new file mode 100644 index 0000000000..d128d18e3a --- /dev/null +++ b/dimos/manipulation/planning/groups/utils.py @@ -0,0 +1,143 @@ +# 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. + +"""Shared helpers for planning-group selectors and joint-state projection.""" + +from collections.abc import Mapping, Sequence + +from dimos.manipulation.planning.groups.identifiers import ( + assert_global_joint_names, + assert_local_joint_names, + is_global_joint_name, +) +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + LocalModelJointName, + PlanningGroupID, +) +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +def planning_group_id_from_selector(selector: PlanningGroupID | PlanningGroup) -> PlanningGroupID: + """Return the planning-group ID represented by a selector.""" + if isinstance(selector, PlanningGroup): + return selector.id + return selector + + +def matching_global_joint_name( + positions_by_name: Mapping[str, float], local_joint_name: LocalModelJointName +) -> GlobalJointName | None: + """Find the unique global joint name ending with a local joint name.""" + suffix = f"/{local_joint_name}" + matches = [name for name in positions_by_name if name.endswith(suffix)] + if len(matches) == 1: + return matches[0] + logger.warning( + f"Expected exactly one global joint ending with local joint name " + f"'{local_joint_name}', found {len(matches)} matches: {matches}" + ) + return None + + +def filter_joint_state_to_selected_joints( + joint_state: JointState, + global_joint_names: Sequence[GlobalJointName], + local_joint_names: Sequence[LocalModelJointName] = (), +) -> JointState: + """Project a joint state to selected global joints. + + Values are looked up by global name first. When ``local_joint_names`` is + provided, each corresponding local name is used as a fallback. + """ + if local_joint_names and len(global_joint_names) != len(local_joint_names): + raise ValueError("Global and local selected joint lists must have the same length") + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + selected_positions: list[float] = [] + missing: list[str] = [] + for index, global_name in enumerate(global_joint_names): + if global_name in positions_by_name: + selected_positions.append(float(positions_by_name[global_name])) + continue + if local_joint_names: + local_name = local_joint_names[index] + if local_name in positions_by_name: + selected_positions.append(float(positions_by_name[local_name])) + continue + missing.append(global_name) + + if missing: + raise ValueError(f"Joint state is missing selected joints: {missing}") + + return JointState({"name": list(global_joint_names), "position": selected_positions}) + + +def joint_target_to_global_names( + group: PlanningGroup, + target: JointState, +) -> JointState: + """Convert a group joint target to global joint names in group order. + + Named targets may use either the public global planning names or the + robot-local model names used by legacy robot-scoped callers, but the two + namespaces must not be mixed in one target. + """ + if not target.name: + if len(target.position) != len(group.joint_names): + raise ValueError( + f"Target for '{group.id}' has {len(target.position)} positions, " + f"expected {len(group.joint_names)}" + ) + return JointState(name=list(group.joint_names), position=list(target.position)) + + if len(target.name) != len(target.position): + raise ValueError( + f"Target for '{group.id}' has {len(target.name)} names but " + f"{len(target.position)} positions" + ) + + target_names = list(target.name) + global_flags = [is_global_joint_name(name) for name in target_names] + if any(global_flags) and not all(global_flags): + raise ValueError( + f"Target for '{group.id}' mixes global and local joint names: {target_names}" + ) + + if all(global_flags): + assert_global_joint_names(target_names) + expected_names = group.joint_names + else: + assert_local_joint_names(target_names) + expected_names = group.local_joint_names + + positions_by_name = dict(zip(target_names, target.position, strict=True)) + global_positions: list[float] = [] + missing: list[str] = [] + for expected_name in expected_names: + if expected_name in positions_by_name: + global_positions.append(positions_by_name[expected_name]) + else: + missing.append(expected_name) + if missing: + raise ValueError(f"Target for '{group.id}' is missing joints: {missing}") + + extra = set(target_names) - set(expected_names) + if extra: + raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") + return JointState(name=list(group.joint_names), position=global_positions) diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index 9bfb1abeff..920f408e22 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -26,6 +26,7 @@ from pytest_mock import MockerFixture from dimos.manipulation.planning.factory import create_kinematics +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig import dimos.manipulation.planning.kinematics.pink_ik as pink_ik from dimos.manipulation.planning.kinematics.pink_ik import ( @@ -169,8 +170,15 @@ def _robot_config() -> RobotModelConfig: model_path=Path("/tmp/fake.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), joint_names=["joint_a", "joint_b", "joint_c"], - end_effector_link="tool", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint_a", "joint_b", "joint_c"), + base_link="base", + tip_link="tool", + ) + ], ) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index df75d23d03..355c1e6712 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -18,6 +18,7 @@ from typing import Any from dimos.manipulation.planning import factory as planning_factory +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.monitor import world_monitor as world_monitor_module from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import PlanningSceneInfo @@ -155,8 +156,12 @@ def _robot_config() -> RobotModelConfig: model_path=Path("/tmp/arm.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0, 0, 0, 1])), joint_names=["j1", "j2"], - end_effector_link="ee", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], ) diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index 74dc3bd69b..1ef40d103b 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -21,6 +21,11 @@ from pydantic import Field from dimos.core.module import ModuleConfig +from dimos.manipulation.planning.groups.identifiers import ( + assert_local_joint_names, + assert_valid_robot_name, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -30,10 +35,13 @@ class RobotModelConfig(ModuleConfig): Attributes: name: Human-readable robot name model_path: Path to robot model file (.urdf, .xacro, or .xml/MJCF) - base_pose: Pose of robot base in world frame (position + orientation) - joint_names: Ordered list of controlled joint names (in URDF namespace) - end_effector_link: Name of the end-effector link for FK/IK - base_link: Name of the base link (default: "base_link") + srdf_path: Optional path to SRDF file containing planning group definitions + base_pose: Placement transform. This is the canonical world placement for + robot instances. + joint_names: Ordered list of controllable joints in the local model + namespace. This is not a planning group. + base_link: Robot-scoped link that base_pose places in the world and + current backends use for weld/placement. package_paths: Dict mapping package names to filesystem Paths joint_limits_lower: Lower joint limits (radians) joint_limits_upper: Upper joint limits (radians) @@ -45,19 +53,20 @@ class RobotModelConfig(ModuleConfig): links may legitimately overlap (e.g., mimic joints). max_velocity: Maximum joint velocity for trajectory generation (rad/s) max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) - joint_name_mapping: Maps coordinator joint names to URDF joint names. - Example: {"left/joint1": "joint1"} means coordinator's "left/joint1" - corresponds to URDF's "joint1". If empty, names are assumed to match. + joint_name_mapping: Maps coordinator joint names to local model joint names. + This is retained for current coordinator/monitor integrations while planning + APIs move toward globally scoped joint names. coordinator_task_name: Task name for executing trajectories via coordinator RPC. If set, trajectories can be executed via execute_trajectory() RPC. """ name: str model_path: Path - base_pose: PoseStamped + srdf_path: Path | None = None + base_pose: PoseStamped = Field(default_factory=PoseStamped) joint_names: list[str] - end_effector_link: str base_link: str = "base_link" + planning_groups: list[PlanningGroupDefinition] = Field(default_factory=list) package_paths: dict[str, Path] = Field(default_factory=dict) joint_limits_lower: list[float] | None = None joint_limits_upper: list[float] | None = None @@ -79,14 +88,44 @@ class RobotModelConfig(ModuleConfig): # Pre-grasp offset distance in meters (along approach direction) pre_grasp_offset: float = 0.10 + def model_post_init(self, __context: object) -> None: + """Validate delimiter-based naming constraints.""" + assert_valid_robot_name(self.name) + assert_local_joint_names(self.joint_names) + + @property + def end_effector_link(self) -> str: + """Compatibility pose target frame derived from planning groups. + + Current world, IK, and visualization layers still ask robot configs for + one end-effector link. The planning-group model stores that frame as a + group ``tip_link``; this shim keeps those layers working until they are + migrated to explicit planning-group IDs. + """ + pose_tip_links = [ + group.tip_link for group in self.planning_groups if group.tip_link is not None + ] + if not pose_tip_links: + raise ValueError( + f"RobotModelConfig '{self.name}' has no pose-target planning group; " + "define PlanningGroupDefinition.tip_link" + ) + unique_tip_links = list(dict.fromkeys(pose_tip_links)) + if len(unique_tip_links) > 1: + raise ValueError( + f"RobotModelConfig '{self.name}' has multiple pose-target planning groups; " + "use an explicit planning group ID" + ) + return unique_tip_links[0] + def get_urdf_joint_name(self, coordinator_name: str) -> str: - """Translate coordinator joint name to URDF joint name.""" + """Translate coordinator joint name to local model joint name.""" return self.joint_name_mapping.get(coordinator_name, coordinator_name) def get_coordinator_joint_name(self, urdf_name: str) -> str: - """Translate URDF joint name to coordinator joint name.""" - for coord_name, u_name in self.joint_name_mapping.items(): - if u_name == urdf_name: + """Translate local model joint name to coordinator joint name.""" + for coord_name, model_name in self.joint_name_mapping.items(): + if model_name == urdf_name: return coord_name return urdf_name @@ -94,4 +133,4 @@ def get_coordinator_joint_names(self) -> list[str]: """Get joint names in coordinator namespace.""" if not self.joint_name_mapping: return self.joint_names - return [self.get_coordinator_joint_name(j) for j in self.joint_names] + return [self.get_coordinator_joint_name(joint_name) for joint_name in self.joint_names] diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index d412e9f766..e746bf2bc2 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -41,6 +41,15 @@ WorldRobotID: TypeAlias = str """Internal Drake world robot ID""" +PlanningGroupID: TypeAlias = str +"""Public planning group ID of the form {robot_name}/{group_name}.""" + +LocalModelJointName: TypeAlias = str +"""Joint name as it appears in URDF/SRDF before world binding.""" + +GlobalJointName: TypeAlias = str +"""Public joint name of the form {robot_name}/{local_joint_name}.""" + JointPath: TypeAlias = "list[JointState]" """List of joint states forming a path (each waypoint has names + positions)""" diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index f8e914e3b3..e0958afb8c 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -30,6 +30,7 @@ ManipulationModule, ManipulationState, ) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -62,8 +63,15 @@ def _get_xarm7_config() -> RobotModelConfig: model_path=desc_path / "urdf/xarm_device.urdf.xacro", base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], - end_effector_link="link7", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"), + base_link="link_base", + tip_link="link7", + ) + ], package_paths={"xarm_description": desc_path}, xacro_args={"dof": "7", "limited": "true"}, auto_convert_meshes=True, diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index bc12bfe994..c4adaf6212 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -28,6 +28,7 @@ ManipulationModuleConfig, ManipulationState, ) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor from dimos.manipulation.planning.spec.config import RobotModelConfig @@ -51,8 +52,15 @@ def robot_config(): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], - end_effector_link="link_tcp", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], max_velocity=1.0, max_acceleration=2.0, coordinator_task_name="traj_arm", @@ -67,8 +75,15 @@ def robot_config_with_mapping(): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], - end_effector_link="link_tcp", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], joint_name_mapping={ "left/joint1": "joint1", "left/joint2": "joint2", @@ -415,7 +430,11 @@ def test_execute_requires_task_name(self): model_path=Path("/path"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["j1"], - end_effector_link="ee", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1",), base_link="base_link", tip_link="ee" + ) + ], ) module._robots = {"arm": ("id", config_no_task, MagicMock())} module._planned_trajectories = {"arm": MagicMock()} @@ -598,8 +617,12 @@ def test_multi_robot_splits_correctly(self): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["j1", "j2"], - end_effector_link="ee", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, coordinator_task_name="traj_left", ) @@ -608,8 +631,12 @@ def test_multi_robot_splits_correctly(self): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["j1", "j2"], - end_effector_link="ee", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, coordinator_task_name="traj_right", ) diff --git a/dimos/manipulation/test_planning_factory.py b/dimos/manipulation/test_planning_factory.py index e695d9579f..99771f9fea 100644 --- a/dimos/manipulation/test_planning_factory.py +++ b/dimos/manipulation/test_planning_factory.py @@ -32,6 +32,7 @@ create_world, validate_backend_combination, ) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import JacobianKinematicsConfig from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner @@ -64,7 +65,14 @@ def robot_config() -> RobotModelConfig: model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], - end_effector_link="tcp", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2"), + base_link="base_link", + tip_link="tcp", + ) + ], coordinator_task_name="traj_arm", ) diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 9457c4264c..65e4000a03 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -25,6 +25,7 @@ import numpy as np import pytest +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import Obstacle @@ -263,7 +264,14 @@ def robot_config(tmp_path: Path) -> RobotModelConfig: model_path=model_path, base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], - end_effector_link="tcp", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2"), + base_link="base", + tip_link="tcp", + ) + ], joint_limits_lower=[-1.0, -2.0], joint_limits_upper=[1.0, 2.0], ) diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index 4d49f019cf..0ad698748d 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -24,6 +24,7 @@ import pytest from dimos.manipulation.manipulation_module import ManipulationModuleConfig +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( JointPath, @@ -78,7 +79,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: model_path=Path("fake.urdf"), base_pose=PoseStamped(), joint_names=[], - end_effector_link="ee_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=(), base_link="base_link", tip_link="ee_link" + ) + ], ) def get_joint_limits( diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index 19f1f2da60..63891bad25 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -21,6 +21,7 @@ pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import PlanningSceneInfo from dimos.manipulation.visualization.viser import ( @@ -63,7 +64,11 @@ def fake_robot_config(name: str) -> RobotModelConfig: model_path=Path(f"{name}.urdf"), base_pose=PoseStamped(), joint_names=[], - end_effector_link="ee_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=(), base_link="base_link", tip_link="ee_link" + ) + ], ) diff --git a/dimos/robot/manipulators/_modeling.py b/dimos/robot/manipulators/_modeling.py index 852a484b63..d9b528d0d5 100644 --- a/dimos/robot/manipulators/_modeling.py +++ b/dimos/robot/manipulators/_modeling.py @@ -16,6 +16,7 @@ from __future__ import annotations +import math from typing import TypeAlias from dimos.manipulation.planning.spec.models import RobotName @@ -31,10 +32,16 @@ JointNameMapping: TypeAlias = dict[CoordinatorJointName, UrdfJointName] -def base_pose(x: float = 0.0, y: float = 0.0, z: float = 0.0) -> PoseStamped: +def base_pose( + x: float = 0.0, + y: float = 0.0, + z: float = 0.0, + pitch: float = 0.0, +) -> PoseStamped: + half_pitch = pitch / 2.0 return PoseStamped( position=Vector3(x=x, y=y, z=z), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + orientation=Quaternion([0.0, math.sin(half_pitch), 0.0, math.cos(half_pitch)]), ) diff --git a/dimos/robot/manipulators/a750/config.py b/dimos/robot/manipulators/a750/config.py index 32bc57409a..4711c2f537 100644 --- a/dimos/robot/manipulators/a750/config.py +++ b/dimos/robot/manipulators/a750/config.py @@ -21,6 +21,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -102,13 +103,21 @@ def make_a750_model_config( coordinator_task_name: str | None = None, ) -> RobotModelConfig: dof = 6 + local_joint_names = joint_names(dof) return RobotModelConfig( name=name, model_path=A750_MODEL_PATH, base_pose=base_pose(), - joint_names=joint_names(dof), - end_effector_link="gripper_base", + joint_names=local_joint_names, base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="base_link", + tip_link="gripper_base", + ) + ], package_paths=A750_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=A750_GRIPPER_COLLISION_EXCLUSIONS, diff --git a/dimos/robot/manipulators/openarm/config.py b/dimos/robot/manipulators/openarm/config.py index 5306d408bf..45b935b336 100644 --- a/dimos/robot/manipulators/openarm/config.py +++ b/dimos/robot/manipulators/openarm/config.py @@ -20,6 +20,7 @@ from typing import Any from dimos.control.components import HardwareComponent, HardwareType +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import base_pose from dimos.utils.data import LfsPath @@ -81,13 +82,21 @@ def openarm_hardware( def openarm_model_config(side: str, name: str | None = None) -> RobotModelConfig: validate_side(side) resolved_name = name or f"{side}_arm" + local_joint_names = openarm_joints(side) return RobotModelConfig( name=resolved_name, model_path=OPENARM_LEFT_MODEL if side == "left" else OPENARM_RIGHT_MODEL, base_pose=base_pose(), - joint_names=openarm_joints(side), - end_effector_link=f"openarm_{side}_link7", + joint_names=local_joint_names, base_link="openarm_body_link0", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="openarm_body_link0", + tip_link=f"openarm_{side}_link7", + ) + ], package_paths=OPENARM_PACKAGE_PATHS, collision_exclusion_pairs=OPENARM_COLLISION_EXCLUSIONS, auto_convert_meshes=True, @@ -112,13 +121,21 @@ def openarm_single_hardware( def openarm_single_model_config() -> RobotModelConfig: + local_joint_names = openarm_joints("left") return RobotModelConfig( name="arm", model_path=OPENARM_V10_FK_MODEL, base_pose=base_pose(), - joint_names=openarm_joints("left"), - end_effector_link="openarm_left_link7", + joint_names=local_joint_names, base_link="openarm_body_link0", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="openarm_body_link0", + tip_link="openarm_left_link7", + ) + ], package_paths=OPENARM_PACKAGE_PATHS, auto_convert_meshes=True, max_velocity=0.5, diff --git a/dimos/robot/manipulators/piper/config.py b/dimos/robot/manipulators/piper/config.py index f99e9589a5..68772166ad 100644 --- a/dimos/robot/manipulators/piper/config.py +++ b/dimos/robot/manipulators/piper/config.py @@ -20,6 +20,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -110,13 +111,21 @@ def make_piper_model_config( home_joints: list[float] | None = None, ) -> RobotModelConfig: dof = 6 + local_joint_names = joint_names(dof) return RobotModelConfig( name=name, model_path=PIPER_MODEL_PATH, base_pose=base_pose(), - joint_names=joint_names(dof), - end_effector_link="gripper_base", + joint_names=local_joint_names, base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="base_link", + tip_link="gripper_base", + ) + ], package_paths=PIPER_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=PIPER_GRIPPER_COLLISION_EXCLUSIONS, diff --git a/dimos/robot/manipulators/xarm/config.py b/dimos/robot/manipulators/xarm/config.py index d918ee94a8..3f6beb7afc 100644 --- a/dimos/robot/manipulators/xarm/config.py +++ b/dimos/robot/manipulators/xarm/config.py @@ -21,6 +21,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -164,19 +165,28 @@ def make_xarm_model_config( xacro_args = { "dof": str(dof), "limited": "true", - "attach_xyz": f"{x_offset} {y_offset} {z_offset}", - "attach_rpy": f"0 {pitch} 0", + "attach_xyz": "0 0 0", + "attach_rpy": "0 0 0", } if add_gripper: xacro_args["add_gripper"] = "true" + local_joint_names = joint_names(dof) + tip_link = "link_tcp" if add_gripper else f"link{dof}" return RobotModelConfig( name=name, model_path=XARM_MODEL_PATH, - base_pose=base_pose(x_offset, y_offset, z_offset), - joint_names=joint_names(dof), - end_effector_link="link_tcp" if add_gripper else f"link{dof}", + base_pose=base_pose(x_offset, y_offset, z_offset, pitch), + joint_names=local_joint_names, base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="link_base", + tip_link=tip_link, + ) + ], package_paths=XARM_PACKAGE_PATHS, xacro_args=xacro_args, auto_convert_meshes=True, diff --git a/docs/capabilities/manipulation/adding_a_custom_arm.md b/docs/capabilities/manipulation/adding_a_custom_arm.md index aaa075ee49..b173ff22d1 100644 --- a/docs/capabilities/manipulation/adding_a_custom_arm.md +++ b/docs/capabilities/manipulation/adding_a_custom_arm.md @@ -483,6 +483,7 @@ Place your URDF/xacro files under LFS data so they can be resolved via `LfsPath` from dimos.utils.data import LfsPath from dimos.manipulation.manipulation_module import manipulation_module from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.manipulation.planning.spec.models import PlanningGroupDefinition from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -506,7 +507,6 @@ def _make_base_pose(x=0.0, y=0.0, z=0.0) -> PoseStamped: def _make_yourarm_config( name: str = "arm", y_offset: float = 0.0, - joint_prefix: str = "", coordinator_task: str | None = None, ) -> RobotModelConfig: """Create YourArm robot config for planning. @@ -514,27 +514,31 @@ def _make_yourarm_config( Args: name: Robot name in the Drake planning world. y_offset: Y-axis offset for multi-arm setups. - joint_prefix: Prefix for joint name mapping to coordinator namespace. coordinator_task: Coordinator task name for trajectory execution via RPC. """ # These must match the joint names in your URDF joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] - joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} return RobotModelConfig( name=name, model_path=_YOURARM_URDF_PATH, - base_pose=_make_base_pose(y=y_offset), joint_names=joint_names, - end_effector_link="link6", # Last link in your URDF's kinematic chain - base_link="base_link", # Root link of your URDF + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(joint_names), + base_link="base_link", + tip_link="link6", + ) + ], + base_pose=_make_base_pose(y=y_offset), # world -> base_link placement + base_link="base_link", # Robot-scoped placement/weld/strip link package_paths={"yourarm_description": _YOURARM_PACKAGE_PATH}, xacro_args={}, # Xacro arguments if using .xacro files collision_exclusion_pairs=[], # Pairs of links that can touch (e.g., gripper fingers) auto_convert_meshes=True, # Convert DAE/STL meshes for Drake max_velocity=1.0, # Max velocity scaling factor max_acceleration=2.0, # Max acceleration scaling factor - joint_name_mapping=joint_mapping, coordinator_task_name=coordinator_task, ) ``` @@ -546,7 +550,7 @@ Add this to your `dimos/robot/yourarm/blueprints.py` alongside the coordinator b ```python skip yourarm_planner = manipulation_module( - robots=[_make_yourarm_config("arm", joint_prefix="arm_", coordinator_task="traj_arm")], + robots=[_make_yourarm_config("arm", coordinator_task="traj_arm")], planning_timeout=10.0, visualization={"backend": "meshcat"}, ) @@ -560,14 +564,23 @@ yourarm_planner = manipulation_module( | Field | Description | |-------|-------------| | `model_path` | Path to `.urdf` or `.xacro` file | -| `joint_names` | Ordered list of controlled joints (must match URDF) | -| `end_effector_link` | Link to use as the end-effector for IK | -| `base_link` | Root link of the robot model | +| `joint_names` | Ordered controllable local model joint set (must match URDF); not itself a planning group | +| `planning_groups` / `srdf_path` | Explicit planning groups or SRDF source; direct `RobotModelConfig(...)` helpers should pass explicit groups, while shared config helpers can discover groups from SRDF/fallback | +| `base_pose` / `base_link` | Optional robot placement: `base_pose` places `base_link` in the world for weld/strip behavior | | `package_paths` | Maps `package://` URIs to filesystem paths (for xacro) | -| `joint_name_mapping` | Maps coordinator names (e.g., `"arm_joint1"`) to URDF names (e.g., `"joint1"`) | | `coordinator_task_name` | Must match the `TaskConfig.name` in your coordinator blueprint | | `collision_exclusion_pairs` | List of `(link_a, link_b)` tuples for links that may legitimately touch (e.g., gripper fingers) | +Coordinator-facing joint states and trajectories use global joint names derived +mechanically as `{robot_name}/{local_joint_name}` (for example, `arm/joint1`). +Keep hardware-native name translation inside the hardware adapter; manipulation +planning config uses local model joint names. + +Planning-group `base_link`/`tip_link` values define kinematic chains and pose +target frames. `base_link` is only the robot-scoped link placed by +`base_pose`; do not use it as a substitute for planning-group chain metadata. +See [Planning Groups](/docs/capabilities/manipulation/planning_groups.md). + ## Step 5: Register Blueprints The blueprint registry in `dimos/robot/all_blueprints.py` is **auto-generated** by scanning the codebase for blueprint declarations. After adding your blueprints: diff --git a/docs/capabilities/manipulation/planning_groups.md b/docs/capabilities/manipulation/planning_groups.md new file mode 100644 index 0000000000..043899fd50 --- /dev/null +++ b/docs/capabilities/manipulation/planning_groups.md @@ -0,0 +1,161 @@ +# Manipulation Planning Groups + +Planning groups are named, selectable kinematic chains used by manipulation +planning. They let APIs target a specific part of a robot, such as an arm or +torso, without confusing that group with the robot's hardware identity. + +## Concepts + +| Concept | Meaning | +|---------|---------| +| Robot name | The configured robot ID in `RobotModelConfig.name`. | +| Planning group | A named serial chain of controllable joints on one robot. | +| Planning group ID | Stable API ID in the form `{robot_name}/{group_name}`. | +| Local joint name | Joint name inside a robot model, such as `joint1`. | +| Global joint name | Boundary-level joint name in the form `{robot_name}/{local_joint_name}`. | +| Generated plan | Planning artifact containing selected group IDs and one synchronized global-joint path. | +| Auxiliary group | A selected group that contributes free DOFs to a pose plan without receiving its own pose target. | + +Local URDF/SRDF joint names stay inside robot-scoped configuration, model +parsing, and backend internals. Flat planning states and generated plan paths +use global joint names so multiple robots can safely share local names such as +`joint1`. + +## Discovering planning groups + +Planning-group config helpers can discover groups in this order: + +1. Explicit `planning_groups` on the robot model config. +2. Explicit `srdf_path` on the robot model config. +3. Conservative SRDF auto-discovery near the model path, with a warning. +4. Fallback generation of one `{robot_name}/manipulator` group when the + configured controllable joints form exactly one unambiguous serial chain. +5. Error if no SRDF or fallback chain can provide a single valid group. + +Direct `RobotModelConfig(...)` construction does not run discovery or synthesize +groups in `model_post_init`; callers must pass explicit `planning_groups` there. + +Supported SRDF group forms: + +```xml + + + +``` + +```xml + + + + + +``` + +Unsupported SRDF forms are skipped with warnings: link groups, nested group +references, mixed group declarations, branching or non-serial groups, and SRDF +`` metadata. A chain group's `tip_link` is its pose target frame. +An ordered joint-list group can be pose-targeted only when DimOS can validate a +unique serial target frame. + +## Fallback behavior + +When no SRDF or explicit group config is available, +fallback uses `RobotModelConfig.joint_names` as the candidate controllable set. +This field is the robot's ordered local model joint set, not an implicit +planning group. + +Fallback succeeds only when those joints form one unambiguous serial chain. It +allows prismatic joints in the middle of the chain and strips only terminal tip +prismatic joints, which usually represent gripper fingers. The generated group +name is always `manipulator`. + +## Current APIs + +Use `list_planning_groups()` to discover group IDs and capabilities before +planning: + +```python skip +groups = manip.list_planning_groups() +pose_groups = [group for group in groups if group.has_pose_target] +group_id = pose_groups[0].id +``` + +Joint-space planning targets group IDs and uses local joint names inside each +target `JointState`: + +```python skip +ok = manip.plan_to_joint_targets( + { + "left_arm/manipulator": JointState( + name=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"], + position=[0.0, -0.4, 0.2, 0.0, 0.3, 0.0], + ) + } +) +``` + +Pose planning targets pose-capable group IDs. Add auxiliary groups when another +chain should participate as free DOFs but does not have its own pose target: + +```python skip +ok = manip.plan_to_pose_targets( + {"left_arm/manipulator": target_pose}, + auxiliary_groups=["torso/manipulator"], +) +``` + +After a successful planning call, preview and execution use the module's current +stored plan: + +```python skip +manip.preview_plan() +manip.execute_plan() +``` + +Callers that already hold a `GeneratedPlan` may pass it explicitly: + +```python skip +manip.preview_plan(plan) +manip.execute_plan(plan) +``` + +For robot-scoped compatibility APIs, unnamed joint vectors are interpreted in +the robot model's configured joint order. If names are provided, they must be +local model joint names: no global names, missing joints, extra joints, or +partial joint sets. + +## Generated plans and execution + +A `GeneratedPlan` stores: + +- selected planning group IDs; +- a single synchronized path of `JointState` waypoints keyed by global joint + names; +- status, timing, path length, iteration count, and message metadata. + +Preview and execution project this path lazily. Preview sends projected joint +paths to the world monitor. Execution splits the path by affected trajectory +task, orders each trajectory by the robot's configured local joint order, writes +global joint names at the coordinator boundary, and invokes each trajectory +controller. Controllers remain planning-group agnostic. + +Multi-task dispatch is not atomic: if one trajectory task accepts and a later +task rejects, DimOS reports the rejection but does not roll back the accepted +task. + +## Robot placement config + +`RobotModelConfig.base_pose` and `RobotModelConfig.base_link` describe robot +placement: `base_pose` places `base_link` in the world and current backends +use that link for weld/placement and optional model-authored world-joint +stripping. This is robot placement metadata, not planning-chain metadata. + +Planning-group `base_link` and `tip_link` values are the only source for chain +bases and pose target frames. Robot-scoped end-effector config is no longer +supported; robot-level EE helper APIs are wrappers over a unique pose-targetable +planning group and should use explicit group APIs when multiple pose groups +exist. + +Robot placement can be encoded either in model assets or in `base_pose`, +depending on the blueprint. `joint_names` remains supported and should describe +the ordered controllable local model joint set. diff --git a/openspec/changes/planning-groups-foundation/proposal.md b/openspec/changes/planning-groups-foundation/proposal.md index c1643bbc5e..c6c7261bc1 100644 --- a/openspec/changes/planning-groups-foundation/proposal.md +++ b/openspec/changes/planning-groups-foundation/proposal.md @@ -8,7 +8,7 @@ This change extracts PR 1 from the reference branch/worktree `cc/spec/movegroup` - Add planning-group identifiers, definitions, runtime models, registry helpers, and joint-name conversion utilities. - Update planning spec/config models so robot-scoped `end_effector_link` is removed and pose target frames are represented by planning-group `tip_link` values. -- Add high-level `RobotConfig` conversion/discovery support for planning groups. +- Add planning-group conversion/discovery support through manipulation planning config models. - Update manipulator robot configs to declare or discover explicit planning groups. - Add foundation-level unit tests and minimal concept docs. @@ -24,5 +24,5 @@ This change extracts PR 1 from the reference branch/worktree `cc/spec/movegroup` - Base branch: `main`. - Reference implementation: branch/worktree `cc/spec/movegroup`; stabilize to a commit before assigning agents. -- Primary files: `dimos/manipulation/planning/groups/*`, `dimos/manipulation/planning/spec/*`, `dimos/robot/config.py`, `dimos/robot/manipulators/*/config.py`, foundation tests, and planning-group docs. +- Primary files: `dimos/manipulation/planning/groups/*`, `dimos/manipulation/planning/spec/*`, `dimos/robot/manipulators/*/config.py`, `dimos/manipulation/planning/groups/test_planning_groups.py`, foundation tests, and planning-group docs. - Out of scope: world backends, IK/RRT algorithms, `ManipulationModule`, Viser UI, and control changes. diff --git a/openspec/changes/planning-groups-foundation/tasks.md b/openspec/changes/planning-groups-foundation/tasks.md index 1e9ef2d761..0b625f1e79 100644 --- a/openspec/changes/planning-groups-foundation/tasks.md +++ b/openspec/changes/planning-groups-foundation/tasks.md @@ -1,18 +1,18 @@ ## 1. Extract foundation files -- [ ] 1.1 Start from `main` and use `cc/spec/movegroup` as the reference implementation. -- [ ] 1.2 Bring over `dimos/manipulation/planning/groups/*`. -- [ ] 1.3 Bring over required planning spec/config changes only. -- [ ] 1.4 Bring over `dimos/robot/config.py` and manipulator config group declarations. +- [x] 1.1 Start from `main` and use `cc/spec/movegroup` as the reference implementation. +- [x] 1.2 Bring over `dimos/manipulation/planning/groups/*`. +- [x] 1.3 Bring over required planning spec/config changes only. +- [x] 1.4 Bring over manipulator config group declarations without `dimos/robot/config.py`. ## 2. Tests and docs -- [ ] 2.1 Bring over `dimos/manipulation/planning/test_planning_groups.py` and `test_planning_group_utils.py`. -- [ ] 2.2 Bring over `dimos/robot/test_config.py`. -- [ ] 2.3 Bring over only foundation docs needed to explain planning groups and custom-arm config. +- [x] 2.1 Bring over `dimos/manipulation/planning/groups/test_planning_groups.py`. +- [x] 2.2 Remove stale `dimos/robot/test_config.py` scope from this PR. +- [x] 2.3 Bring over only foundation docs needed to explain planning groups and custom-arm config. ## 3. Validation -- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/test_planning_groups.py dimos/manipulation/planning/test_planning_group_utils.py dimos/robot/test_config.py -q`. -- [ ] 3.2 Run targeted mypy on changed production files. -- [ ] 3.3 Confirm no files from world/IK/RRT/module/Viser/control were pulled in accidentally. +- [x] 3.1 Run targeted planning-group tests under `dimos/manipulation/planning/groups/test_planning_groups.py`. +- [x] 3.2 Run targeted mypy on changed production files. +- [x] 3.3 Confirm no files from world/IK/RRT/module/Viser/control were pulled in accidentally. From c7eb5d93497aa9364f48dc0c86c74b0409b4f827 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 16:47:31 -0700 Subject: [PATCH 03/29] spec: archive 1 --- .../planning-groups-foundation/.openspec.yaml | 2 -- .../planning-groups-foundation/README.md | 3 -- .../planning-groups-foundation/design.md | 29 ------------------- .../planning-groups-foundation/proposal.md | 28 ------------------ .../planning-groups-foundation/tasks.md | 18 ------------ .../manipulation-planning-config/spec.md | 4 ++- .../specs/planning-group-foundation/spec.md | 4 ++- 7 files changed, 6 insertions(+), 82 deletions(-) delete mode 100644 openspec/changes/planning-groups-foundation/.openspec.yaml delete mode 100644 openspec/changes/planning-groups-foundation/README.md delete mode 100644 openspec/changes/planning-groups-foundation/design.md delete mode 100644 openspec/changes/planning-groups-foundation/proposal.md delete mode 100644 openspec/changes/planning-groups-foundation/tasks.md rename openspec/{changes/planning-groups-foundation => }/specs/manipulation-planning-config/spec.md (90%) rename openspec/{changes/planning-groups-foundation => }/specs/planning-group-foundation/spec.md (95%) diff --git a/openspec/changes/planning-groups-foundation/.openspec.yaml b/openspec/changes/planning-groups-foundation/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/planning-groups-foundation/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/planning-groups-foundation/README.md b/openspec/changes/planning-groups-foundation/README.md deleted file mode 100644 index 820f444c27..0000000000 --- a/openspec/changes/planning-groups-foundation/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# planning-groups-foundation - -Extract PR 1 from reference branch cc/spec/movegroup: introduce planning-group foundations, RobotConfig group generation, manipulator group declarations, and foundation docs/tests. diff --git a/openspec/changes/planning-groups-foundation/design.md b/openspec/changes/planning-groups-foundation/design.md deleted file mode 100644 index 287006f709..0000000000 --- a/openspec/changes/planning-groups-foundation/design.md +++ /dev/null @@ -1,29 +0,0 @@ -## Context - -The current branch introduces planning groups across the stack, but the reviewable foundation is the abstraction itself: stable group IDs, local/global joint mapping, group definitions, RobotConfig discovery/generation, and removal of robot-scoped end-effector metadata. - -## Goals / Non-Goals - -**Goals:** -- Define planning groups as the only source for chain base links and pose target tip links. -- Preserve robot-scoped `base_link` as placement/weld/strip metadata, not a planning-chain selector. -- Let robot configs create explicit or discovered planning groups. -- Provide small, direct tests that prove group IDs, joint-name mapping, discovery, and strict config validation. - -**Non-Goals:** -- Do not migrate world backends or planners in this PR. -- Do not migrate `ManipulationModule` public APIs in this PR. -- Do not include visualization or control changes. - -## Decisions - -- Use `PlanningGroupDefinition` for config-time group declarations and `PlanningGroup` for resolved runtime group IDs. -- Use globally scoped group IDs such as `/` at API boundaries. -- Keep `RobotModelConfig.base_link` for robot placement semantics only. -- Reject legacy robot-scoped `end_effector_link`; pose target frames come from group `tip_link` or SRDF discovery. -- Keep this PR compatible enough for later stacked PRs; if a later change is needed for tests, prefer a small compatibility wrapper over pulling later files forward. - -## Risks / Trade-offs - -- Earlier PRs may need temporary compatibility paths because the full branch removes legacy behavior across several layers. -- Discovery behavior must be clear: high-level `RobotConfig.to_robot_model_config()` can discover/fallback groups, while direct `RobotModelConfig(...)` requires explicit groups. diff --git a/openspec/changes/planning-groups-foundation/proposal.md b/openspec/changes/planning-groups-foundation/proposal.md deleted file mode 100644 index c6c7261bc1..0000000000 --- a/openspec/changes/planning-groups-foundation/proposal.md +++ /dev/null @@ -1,28 +0,0 @@ -## Why - -The planning-group refactor is currently bundled in one large branch. Reviewers need a small first PR that establishes the core planning-group abstraction and robot configuration plumbing without also reviewing planner algorithms, module APIs, visualization, or control integration. - -This change extracts PR 1 from the reference branch/worktree `cc/spec/movegroup`. Treat the existing branch as a working reference implementation; do not redesign the feature unless required to keep this slice independently valid. - -## What Changes - -- Add planning-group identifiers, definitions, runtime models, registry helpers, and joint-name conversion utilities. -- Update planning spec/config models so robot-scoped `end_effector_link` is removed and pose target frames are represented by planning-group `tip_link` values. -- Add planning-group conversion/discovery support through manipulation planning config models. -- Update manipulator robot configs to declare or discover explicit planning groups. -- Add foundation-level unit tests and minimal concept docs. - -## Capabilities - -### New Capabilities -- `planning-group-foundation`: Core planning-group data model, discovery, robot config generation, and validation contracts. - -### Modified Capabilities -- `manipulation-planning-config`: Robot model configuration uses planning groups as the source of chain and pose-target metadata. - -## Impact - -- Base branch: `main`. -- Reference implementation: branch/worktree `cc/spec/movegroup`; stabilize to a commit before assigning agents. -- Primary files: `dimos/manipulation/planning/groups/*`, `dimos/manipulation/planning/spec/*`, `dimos/robot/manipulators/*/config.py`, `dimos/manipulation/planning/groups/test_planning_groups.py`, foundation tests, and planning-group docs. -- Out of scope: world backends, IK/RRT algorithms, `ManipulationModule`, Viser UI, and control changes. diff --git a/openspec/changes/planning-groups-foundation/tasks.md b/openspec/changes/planning-groups-foundation/tasks.md deleted file mode 100644 index 0b625f1e79..0000000000 --- a/openspec/changes/planning-groups-foundation/tasks.md +++ /dev/null @@ -1,18 +0,0 @@ -## 1. Extract foundation files - -- [x] 1.1 Start from `main` and use `cc/spec/movegroup` as the reference implementation. -- [x] 1.2 Bring over `dimos/manipulation/planning/groups/*`. -- [x] 1.3 Bring over required planning spec/config changes only. -- [x] 1.4 Bring over manipulator config group declarations without `dimos/robot/config.py`. - -## 2. Tests and docs - -- [x] 2.1 Bring over `dimos/manipulation/planning/groups/test_planning_groups.py`. -- [x] 2.2 Remove stale `dimos/robot/test_config.py` scope from this PR. -- [x] 2.3 Bring over only foundation docs needed to explain planning groups and custom-arm config. - -## 3. Validation - -- [x] 3.1 Run targeted planning-group tests under `dimos/manipulation/planning/groups/test_planning_groups.py`. -- [x] 3.2 Run targeted mypy on changed production files. -- [x] 3.3 Confirm no files from world/IK/RRT/module/Viser/control were pulled in accidentally. diff --git a/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md b/openspec/specs/manipulation-planning-config/spec.md similarity index 90% rename from openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md rename to openspec/specs/manipulation-planning-config/spec.md index ecb76b89b5..bd13e49196 100644 --- a/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md +++ b/openspec/specs/manipulation-planning-config/spec.md @@ -1,4 +1,6 @@ -## MODIFIED Requirements +# manipulation-planning-config Specification + +## Requirements ### Requirement: Robot model configuration must separate placement from planning chains Robot model configuration MUST treat robot-scoped `base_link` as the link placed by `base_pose` and used by backend weld/strip behavior. Planning-chain base links and pose target frames MUST come from planning groups. diff --git a/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md b/openspec/specs/planning-group-foundation/spec.md similarity index 95% rename from openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md rename to openspec/specs/planning-group-foundation/spec.md index 2da6cd9e9e..6124a3bf5c 100644 --- a/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md +++ b/openspec/specs/planning-group-foundation/spec.md @@ -1,4 +1,6 @@ -## ADDED Requirements +# planning-group-foundation Specification + +## Requirements ### Requirement: Planning group definitions must identify kinematic chains The planning configuration MUST represent kinematic chain metadata with planning-group definitions containing group-local joint names, a chain base link, and an optional pose-target tip link. From 94cd3ce72d4b347aa5fc1caa13ff9b24ec38c52c Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 19:31:20 -0700 Subject: [PATCH 04/29] feat: add planning group world queries --- .../planning/monitor/test_world_monitor.py | 150 +++++++++++++++- .../planning/monitor/world_monitor.py | 166 +++++++++++++++++- dimos/manipulation/planning/spec/protocols.py | 9 + .../planning/world/drake_world.py | 145 ++++++++++++--- .../planning/world/roboplan_world.py | 98 +++++++++-- .../world/test_drake_world_planning_groups.py | 153 ++++++++++++++++ dimos/manipulation/test_roboplan_world.py | 90 ++++++++++ .../planning-world-monitor-groups/tasks.md | 21 +-- 8 files changed, 780 insertions(+), 52 deletions(-) create mode 100644 dimos/manipulation/planning/world/test_drake_world_planning_groups.py diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 355c1e6712..aed6d6a9a5 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -17,6 +17,9 @@ from pathlib import Path from typing import Any +import numpy as np +import pytest + from dimos.manipulation.planning import factory as planning_factory from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.monitor import world_monitor as world_monitor_module @@ -25,15 +28,48 @@ 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.sensor_msgs.JointState import JointState + + +class _VectorLike(list[float]): + def tolist(self) -> list[float]: + return list(self) + + +class _FakeStateMonitor: + def __init__(self, positions: list[float], stale: bool = False) -> None: + self._positions = _VectorLike(positions) + self._stale = stale + + def get_current_positions(self) -> _VectorLike: + return self._positions + + def get_current_velocities(self) -> None: + return None + + def is_state_stale(self, max_age: float) -> bool: + _ = max_age + return self._stale + + +class _ScratchContext: + def __enter__(self) -> str: + return "scratch" + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: + return False class FakeWorld: def __init__(self) -> None: - self.calls: list[tuple[str, Any]] = [] + self.calls: list[tuple[Any, ...]] = [] + self.configs: dict[str, RobotModelConfig] = {} def add_robot(self, config): self.calls.append(("add_robot", config)) - return "robot-1" + robot_id = f"robot-{len(self.configs) + 1}" + self.configs[robot_id] = config + return robot_id def get_robot_ids(self): return [] @@ -70,12 +106,14 @@ def get_live_context(self): return None def scratch_context(self): - return self + self.calls.append(("scratch_context", None)) + return _ScratchContext() def sync_from_joint_state(self, robot_id, joint_state) -> None: return None def set_joint_state(self, ctx, robot_id, joint_state) -> None: + self.calls.append(("set_joint_state", ctx, robot_id, joint_state)) return None def get_joint_state(self, ctx, robot_id): @@ -96,12 +134,20 @@ def check_edge_collision_free(self, robot_id, start, end, step_size: float = 0.0 def get_ee_pose(self, ctx, robot_id): return None + def get_group_ee_pose(self, ctx, group_id): + self.calls.append(("get_group_ee_pose", ctx, group_id)) + return PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + def get_link_pose(self, ctx, robot_id, link_name): return [] def get_jacobian(self, ctx, robot_id): return [] + def get_group_jacobian(self, ctx, group_id): + self.calls.append(("get_group_jacobian", ctx, group_id)) + return np.ones((6, 2)) + def get_visualization_url(self): return None @@ -165,6 +211,26 @@ def _robot_config() -> RobotModelConfig: ) +def _robot_config_with_groups(groups: list[PlanningGroupDefinition]) -> RobotModelConfig: + return _robot_config().model_copy(update={"planning_groups": groups}) + + +def _three_joint_reordered_group_config() -> RobotModelConfig: + return _robot_config().model_copy( + update={ + "joint_names": ["j1", "j2", "j3"], + "planning_groups": [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j2", "j1"), + base_link="base", + tip_link="ee", + ) + ], + } + ) + + def test_world_monitor_add_robot_records_scene_without_visualization_probe() -> None: fake_world = FakeWorld() fake_viz = FakeViz() @@ -209,3 +275,81 @@ def test_create_planning_specs_wraps_existing_world(monkeypatch) -> None: assert planning_specs.world_monitor.visualization is None assert planning_specs.kinematics is fake_kinematics assert planning_specs.planner is fake_planner + + +def test_world_monitor_exposes_planning_groups_and_duplicate_names_do_not_mutate() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_robot_config()) + + assert [group.id for group in monitor.planning_groups.list()] == ["arm/manipulator"] + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(_robot_config()) + assert [call[0] for call in fake_world.calls].count("add_robot") == 1 + + +def test_current_group_joint_state_uses_public_names_in_group_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + state = monitor.get_current_group_joint_state("arm/manipulator") + + assert state.name == ["arm/j2", "arm/j1"] + assert state.position == [0.2, 0.1] + + +def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + pose = monitor.get_group_ee_pose( + "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.8, 0.9]}) + ) + jacobian = monitor.get_group_jacobian( + "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.3, 0.4]}) + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.9, 0.8, 0.3] + assert set_calls[1][3].name == ["j1", "j2", "j3"] + assert set_calls[1][3].position == [0.4, 0.3, 0.3] + assert pose.position.x == 1 + assert jacobian.shape == (6, 2) + assert ("get_group_ee_pose", "scratch", "arm/manipulator") in fake_world.calls + assert ("get_group_jacobian", "scratch", "arm/manipulator") in fake_world.calls + + +def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + no_pose_id = monitor.add_robot( + _robot_config_with_groups( + [PlanningGroupDefinition(name="base", joint_names=("j1",), base_link="base")] + ) + ) + with pytest.raises(ValueError, match="no pose-targetable"): + monitor.get_ee_pose(no_pose_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]})) + + fake_world2 = FakeWorld() + monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] + ambiguous_id = monitor2.add_robot( + _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="a", joint_names=("j1",), base_link="base", tip_link="ee1" + ), + PlanningGroupDefinition( + name="b", joint_names=("j2",), base_link="base", tip_link="ee2" + ), + ] + ) + ) + with pytest.raises(ValueError, match="pose-targetable planning groups"): + monitor2.get_jacobian( + ambiguous_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]}) + ) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 5e12568874..3cd54d05c4 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -21,6 +21,8 @@ from typing import TYPE_CHECKING, Any from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +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.spec.models import PlanningSceneInfo @@ -40,6 +42,8 @@ CollisionObjectMessage, JointPath, Obstacle, + PlanningGroupID, + RobotName, WorldRobotID, ) from dimos.msgs.vision_msgs.Detection3D import Detection3D @@ -61,6 +65,8 @@ def __init__( self._lock = threading.RLock() self._robot_joints: dict[WorldRobotID, list[str]] = {} self._robot_configs: dict[WorldRobotID, RobotModelConfig] = {} + self._robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + self._planning_groups = PlanningGroupRegistry() self._state_monitors: dict[WorldRobotID, RobotStateMonitor] = {} self._obstacle_monitor: WorldObstacleMonitor | None = None self._viz_thread: threading.Thread | None = None @@ -72,12 +78,23 @@ def __init__( def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot. Returns robot_id.""" with self._lock: + if config.name in self._robot_ids_by_name: + raise ValueError(f"Robot name '{config.name}' is already registered") + new_groups = PlanningGroupRegistry() + new_groups.add_robot(config) robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names self._robot_configs[robot_id] = config + self._robot_ids_by_name[config.name] = robot_id + self._planning_groups.add_robot(config) logger.info(f"Added robot '{config.name}' as '{robot_id}'") return robot_id + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered public planning groups.""" + return self._planning_groups + def planning_scene_info(self) -> PlanningSceneInfo: """Return a stable metadata snapshot of the initialized planning scene.""" with self._lock: @@ -294,6 +311,116 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: ctx = self._world.get_live_context() return self._world.get_joint_state(ctx, robot_id) + def current_global_joint_state(self, max_age: float = 1.0) -> JointState: + """Return current state for all fresh robots with public global joint names.""" + names: list[str] = [] + positions: list[float] = [] + for robot_name, robot_id in self._robot_ids_by_name.items(): + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + continue + state = self.get_current_joint_state(robot_id) + if state is None: + continue + pos_by_name = dict(zip(state.name, state.position, strict=True)) + for local_name in self._robot_joints.get(robot_id, []): + if local_name in pos_by_name: + names.append(f"{robot_name}/{local_name}") + positions.append(float(pos_by_name[local_name])) + return JointState({"name": names, "position": positions}) + + def get_current_group_joint_state( + self, group_id: PlanningGroupID, max_age: float = 1.0 + ) -> JointState: + """Return current joint state scoped and ordered for one planning group.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + state = self.get_current_joint_state(robot_id) + if state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + return filter_joint_state_to_selected_joints( + state, group.joint_names, group.local_joint_names + ) + + def _current_robot_joint_state_for_group( + self, group_id: PlanningGroupID, max_age: float = 1.0 + ) -> JointState: + """Return current full robot state for a planning group.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + state = self.get_current_joint_state(robot_id) + if state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + return state + + def _full_robot_joint_state_from_input( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + config = self._robot_configs[robot_id] + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + def _joint_state_for_group_query( + self, group_id: PlanningGroupID, joint_state: JointState | None + ) -> JointState: + """Normalize full-robot or group-scoped input to full robot-local state.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if joint_state is None: + return self._current_robot_joint_state_for_group(group_id) + full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) + if full_state is not None: + return full_state + current_state = self._current_robot_joint_state_for_group(group_id) + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names + ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" if robot_id in self._state_monitors: @@ -367,15 +494,23 @@ def get_ee_pose( self, robot_id: WorldRobotID, joint_state: JointState | None = None ) -> PoseStamped: """Get end-effector pose. Uses current state if joint_state is None.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_ee_pose(group_id, joint_state) + + def get_group_ee_pose( + self, group_id: PlanningGroupID, joint_state: JointState | None = None + ) -> PoseStamped: + """Get planning-group tip pose. Uses current group state if joint_state is None.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - # If no state provided, fetch current from state monitor - if joint_state is None: - joint_state = self.get_current_joint_state(robot_id) - - if joint_state is not None: - self._world.set_joint_state(ctx, robot_id, joint_state) + normalized_state = self._joint_state_for_group_query(group_id, joint_state) + self._world.set_joint_state(ctx, robot_id, normalized_state) - return self._world.get_ee_pose(ctx, robot_id) + return self._world.get_group_ee_pose(ctx, group_id) def get_link_pose( self, robot_id: WorldRobotID, link_name: str, joint_state: JointState | None = None @@ -411,9 +546,22 @@ def get_link_pose( def get_jacobian(self, robot_id: WorldRobotID, joint_state: JointState) -> NDArray[np.float64]: """Get 6xN Jacobian matrix.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_jacobian(group_id, joint_state) + + def get_group_jacobian( + self, group_id: PlanningGroupID, joint_state: JointState + ) -> NDArray[np.float64]: + """Get 6xN planning-group Jacobian matrix.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - self._world.set_joint_state(ctx, robot_id, joint_state) - return self._world.get_jacobian(ctx, robot_id) + normalized_state = self._joint_state_for_group_query(group_id, joint_state) + self._world.set_joint_state(ctx, robot_id, normalized_state) + return self._world.get_group_jacobian(ctx, group_id) # Lifecycle diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index c7ee95ee0a..b5fea13a36 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -33,6 +33,7 @@ IKResult, JointPath, Obstacle, + PlanningGroupID, PlanningResult, PlanningSceneInfo, WorldRobotID, @@ -170,6 +171,14 @@ def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian (6 x n_joints).""" ... + def get_group_ee_pose(self, ctx: Any, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + ... + + def get_group_jacobian(self, ctx: Any, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get planning-group Jacobian (6 x n_group_joints).""" + ... + @runtime_checkable class VisualizationSpec(Protocol): diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index ca426ba340..12067f647c 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -25,11 +25,14 @@ import numpy as np +from dimos.manipulation.planning.groups.identifiers import is_global_joint_name +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( JointPath, Obstacle, + PlanningGroupID, PlanningSceneInfo, WorldRobotID, ) @@ -184,6 +187,7 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: # Tracking data self._robots: dict[WorldRobotID, _RobotData] = {} + self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, _ObstacleData] = {} self._robot_counter = 0 self._obstacle_counter = 0 @@ -207,6 +211,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: + if self._planning_groups.groups_for_robot(config.name): + raise ValueError(f"Robot name '{config.name}' is already registered") + new_groups = PlanningGroupRegistry() + new_groups.add_robot(config) + self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -215,9 +224,16 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._validate_joints(config, model_instance) - ee_frame = self._plant.GetBodyByName( - config.end_effector_link, model_instance - ).body_frame() + ee_link = config.base_link + try: + primary_group_id = new_groups.primary_pose_group_id_for_robot(config.name) + except ValueError: + primary_group_id = None + if primary_group_id is not None: + primary_group = new_groups.get(primary_group_id) + if primary_group.tip_link is not None: + ee_link = primary_group.tip_link + ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() base_frame = self._plant.GetBodyByName(config.base_link, model_instance).body_frame() # Preview (yellow ghost) — always a separate instance per robot @@ -235,6 +251,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: base_frame=base_frame, preview_model_instance=preview_model_instance, ) + self._planning_groups.add_robot(config) logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id @@ -317,6 +334,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered planning groups for this world.""" + return self._planning_groups + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -768,8 +790,7 @@ def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) if not self._finalized or self._plant_context is None: return # Silently ignore before finalization - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) with self._lock: self._set_positions_internal(self._plant_context, robot_id, positions) @@ -787,8 +808,7 @@ def set_joint_state( if not self._finalized: raise RuntimeError("World must be finalized first") - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) # Get plant context from diagram context plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) @@ -809,6 +829,51 @@ def _set_positions_internal( self._plant.SetPositions(plant_ctx, full_positions) + def _joint_state_to_q( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> NDArray[np.float64]: + """Normalize unnamed, robot-local, mapped, or global JointState to robot joint order.""" + if robot_id not in self._robots: + raise KeyError(f"Robot '{robot_id}' not found") + robot_data = self._robots[robot_id] + if not joint_state.name: + if len(joint_state.position) != len(robot_data.config.joint_names): + raise ValueError("JointState position length must match configured joint count") + return np.asarray(joint_state.position, dtype=np.float64) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + name_to_pos: dict[str, float] = {} + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in robot_data.config.joint_names: + resolved_name = name + elif name in robot_data.config.joint_name_mapping: + resolved_name = robot_data.config.joint_name_mapping[name] + elif is_global_joint_name(name): + prefix = f"{robot_data.config.name}/" + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in robot_data.config.joint_names: + raise ValueError(f"Unknown global joint name for DrakeWorld: {name}") + else: + resolved_name = robot_data.config.get_urdf_joint_name(name) + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) + missing = [name for name in robot_data.config.joint_names if name not in name_to_pos] + if missing: + raise ValueError(f"JointState missing joints for DrakeWorld: {missing}") + return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) + + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_groups.get(group_id) + matches = [ + rid for rid, data in self._robots.items() if data.config.name == group.robot_name + ] + if not matches: + raise KeyError(f"No robot registered for planning group '{group_id}'") + return matches[0] + def get_joint_state(self, ctx: Context, robot_id: WorldRobotID) -> JointState: """Get robot joint state from given context.""" if not self._finalized: @@ -905,16 +970,28 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose.""" - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) - ee_body = robot_data.ee_frame.body() + ee_body = self._plant.GetBodyByName(group.tip_link, robot_data.model_instance) X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) # Extract position and quaternion from Drake transform @@ -955,30 +1032,58 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float Rows: [vx, vy, vz, wx, wy, wz] (linear, then angular) """ - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get geometric Jacobian (6 x group joints) in group-local order.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + tip_frame = self._plant.GetBodyByName( + group.tip_link, robot_data.model_instance + ).body_frame() # Compute full Jacobian J_full = self._plant.CalcJacobianSpatialVelocity( plant_ctx, JacobianWrtVariable.kQDot, - robot_data.ee_frame, + tip_frame, np.array([0.0, 0.0, 0.0]), # type: ignore[arg-type] # Point on end-effector self._plant.world_frame(), self._plant.world_frame(), ) - # Extract columns for this robot's joints - n_joints = len(robot_data.joint_indices) + # Extract columns for configured controllable joints only. + joint_indices_by_name = dict( + zip(robot_data.config.joint_names, robot_data.joint_indices, strict=True) + ) + missing = [ + joint_name + for joint_name in group.local_joint_names + if joint_name not in joint_indices_by_name + ] + if missing: + raise ValueError( + f"Planning group '{group_id}' references non-controllable joints: {missing}" + ) + group_joint_indices = [joint_indices_by_name[name] for name in group.local_joint_names] + n_joints = len(group_joint_indices) J_robot = np.zeros((6, n_joints)) - for i, joint_idx in enumerate(robot_data.joint_indices): + for i, joint_idx in enumerate(group_joint_indices): J_robot[:, i] = J_full[:, joint_idx] # Reorder rows: Drake uses [angular, linear], we want [linear, angular] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index d252d3fdc9..e224ff583a 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -41,9 +41,16 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc +from dimos.manipulation.planning.groups.identifiers import is_global_joint_name +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus -from dimos.manipulation.planning.spec.models import Obstacle, PlanningResult, WorldRobotID +from dimos.manipulation.planning.spec.models import ( + Obstacle, + PlanningGroupID, + PlanningResult, + WorldRobotID, +) from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -88,6 +95,7 @@ def __init__(self, enable_viz: bool = False, **_: object) -> None: logger.warning("RoboPlanWorld does not currently provide manipulation visualization") self._robots: dict[WorldRobotID, _RoboPlanRobotData] = {} + self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, Obstacle] = {} self._robot_counter = 0 self._finalized = False @@ -104,6 +112,8 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise ValueError("RoboPlanWorld currently supports one robot per Scene") if not Path(config.model_path).exists(): raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") + if self._planning_groups.groups_for_robot(config.name): + raise ValueError(f"Robot name '{config.name}' is already registered") self._validate_robot_config(config) self._robot_counter += 1 @@ -119,6 +129,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._live_context.q_by_robot[robot_id] = np.zeros( len(config.joint_names), dtype=np.float64 ) + self._planning_groups.add_robot(config) logger.info(f"Added RoboPlan robot '{robot_id}' ({config.name})") return robot_id @@ -130,6 +141,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration by ID.""" return self._get_robot(robot_id).config + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered planning groups for this world.""" + return self._planning_groups + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -272,7 +288,17 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - mat = self.get_link_pose(ctx, robot_id, robot.config.end_effector_link) + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose if RoboPlan exposes FK.""" + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + mat = self.get_link_pose(ctx, self._robot_id_for_group(group_id), group.tip_link) pose = matrix_to_pose(mat) return PoseStamped( frame_id="world", @@ -299,17 +325,46 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian( + self, ctx: RoboPlanContext, group_id: PlanningGroupID + ) -> NDArray[np.float64]: + """Get planning-group Jacobian projected to group-local joint order.""" + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_id = self._robot_id_for_group(group_id) + robot = self._get_robot(robot_id) q = ctx.q_by_robot.get(robot_id) if q is None: raise KeyError(f"Robot '{robot_id}' not found in context") scene = self._require_scene() - result = scene.computeFrameJacobian( - self._to_scene_q(robot_id, q), robot.config.end_effector_link, True - ) + result = scene.computeFrameJacobian(self._to_scene_q(robot_id, q), group.tip_link, True) arr = np.asarray(result, dtype=np.float64) if arr.shape[0] != 6: raise ValueError(f"Unexpected RoboPlan Jacobian shape: {arr.shape}; expected 6 x n") - return arr + scene_joint_order = self._query_scene_joint_order(scene, robot.config) + if scene_joint_order is not None and arr.shape[1] == len(scene_joint_order): + missing = [name for name in group.local_joint_names if name not in scene_joint_order] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [scene_joint_order.index(name) for name in group.local_joint_names] + return arr[:, indices] + if arr.shape[1] == len(robot.config.joint_names): + missing = [ + name for name in group.local_joint_names if name not in robot.config.joint_names + ] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [robot.config.joint_names.index(name) for name in group.local_joint_names] + return arr[:, indices] + raise ValueError( + f"Unexpected RoboPlan Jacobian shape: {arr.shape}; cannot project group '{group_id}'" + ) # PlannerSpec for native RoboPlan planning @@ -507,6 +562,15 @@ def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id] + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_groups.get(group_id) + matches = [ + rid for rid, data in self._robots.items() if data.config.name == group.robot_name + ] + if not matches: + raise KeyError(f"No robot registered for planning group '{group_id}'") + return matches[0] + def _joint_state_to_q( self, robot_id: WorldRobotID, joint_state: JointState ) -> NDArray[np.float64]: @@ -515,10 +579,24 @@ def _joint_state_to_q( raise ValueError("JointState position length must match configured joint count") if not joint_state.name: return np.asarray(joint_state.position, dtype=np.float64) - name_to_pos = { - robot.config.get_urdf_joint_name(name): position - for name, position in zip(joint_state.name, joint_state.position, strict=True) - } + name_to_pos: dict[str, float] = {} + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in robot.config.joint_names: + resolved_name = name + elif name in robot.config.joint_name_mapping: + resolved_name = robot.config.joint_name_mapping[name] + elif is_global_joint_name(name): + prefix = f"{robot.config.name}/" + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in robot.config.joint_names: + raise ValueError(f"Unknown global joint name for RoboPlanWorld: {name}") + else: + resolved_name = robot.config.get_urdf_joint_name(name) + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) missing = [name for name in robot.config.joint_names if name not in name_to_pos] if missing: raise ValueError(f"JointState missing joints for RoboPlanWorld: {missing}") diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py new file mode 100644 index 0000000000..841ed27f99 --- /dev/null +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -0,0 +1,153 @@ +# 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. + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +pydrake = pytest.importorskip("pydrake") +_ = pydrake + +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.world.drake_world import DrakeWorld +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +def _write_urdf(path: Path) -> None: + path.write_text( + """ + + + + + + + + + + + + + + + +""" + ) + + +def _config( + path: Path, groups: list[PlanningGroupDefinition], joints: list[str] | None = None +) -> RobotModelConfig: + return RobotModelConfig( + name="arm", + model_path=path, + base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), + joint_names=joints or ["joint1", "joint2"], + base_link="base_link", + planning_groups=groups, + ) + + +def _arm_group( + *joint_names: str, tip_link: str | None = "tool0", name: str = "arm" +) -> PlanningGroupDefinition: + return PlanningGroupDefinition( + name=name, joint_names=joint_names, base_link="base_link", tip_link=tip_link + ) + + +def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")])) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + group_pose = world.get_group_ee_pose(ctx, "arm/arm") + legacy_pose = world.get_ee_pose(ctx, robot_id) + + assert group_pose.position.x == pytest.approx(2.0) + assert legacy_pose.position.x == pytest.approx(group_pose.position.x) + assert world.get_jacobian(ctx, robot_id).shape == (6, 2) + + +def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot( + _config( + urdf, + [ + _arm_group("joint1", "joint2", name="wrist_forward"), + _arm_group("joint2", "joint1", name="wrist_reverse"), + ], + ) + ) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + forward_jacobian = world.get_group_jacobian(ctx, "arm/wrist_forward") + reverse_jacobian = world.get_group_jacobian(ctx, "arm/wrist_reverse") + + assert reverse_jacobian.shape == (6, 2) + np.testing.assert_allclose(reverse_jacobian[:, 0], forward_jacobian[:, 1]) + np.testing.assert_allclose(reverse_jacobian[:, 1], forward_jacobian[:, 0]) + + +def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + no_pose = DrakeWorld() + no_pose_id = no_pose.add_robot(_config(urdf, [_arm_group("joint1", tip_link=None)])) + no_pose.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose.get_ee_pose(no_pose.get_live_context(), no_pose_id) + + ambiguous = DrakeWorld() + ambiguous_id = ambiguous.add_robot( + _config( + urdf, + [ + _arm_group("joint1", tip_link="link1", name="a"), + _arm_group("joint2", tip_link="tool0", name="b"), + ], + ) + ) + ambiguous.finalize() + with pytest.raises(ValueError, match="multiple pose-targetable"): + ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) + + +def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")], joints=["joint1"])) + world.finalize() + + with pytest.raises(ValueError, match="non-controllable"): + world.get_group_jacobian(world.get_live_context(), "arm/arm") diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 65e4000a03..f3879297d5 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -387,6 +387,37 @@ def test_joint_name_mapping_is_applied_to_input_states( assert live_round_trip.position == [0.2, 0.3] +def test_global_joint_names_are_mapped_without_regressing_coordinator_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"arm/j1": "joint1", "arm/j2": "joint2"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/j1", "arm/j2"], position=[0.4, 0.5]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.4, 0.5] + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/joint1", "arm/joint2"], position=[0.2, 0.3]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.2, 0.3] + + +def test_duplicate_resolved_joint_names_fail_clearly( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"alias": "joint1"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(ValueError, match="duplicate joint 'joint1'"): + world.sync_from_joint_state( + robot_id, JointState(name=["joint1", "alias"], position=[0.1, 0.2]) + ) + + def test_obstacle_mutation_updates_scene_and_stored_pose( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: @@ -473,6 +504,65 @@ def test_fk_jacobian_and_explicit_min_distance_unsupported( world.get_min_distance(ctx, robot_id) +def test_group_fk_and_jacobian_use_group_tip_and_local_joint_order( + fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + config = robot_config.model_copy( + update={ + "joint_names": ["joint1", "joint2", "joint3"], + "planning_groups": [ + PlanningGroupDefinition( + name="wrist", + joint_names=("joint3", "joint1"), + base_link="base", + tip_link="wrist_tip", + ) + ], + "joint_limits_lower": [-1.0, -2.0, -3.0], + "joint_limits_upper": [1.0, 2.0, 3.0], + } + ) + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint2", "joint1", "joint3"]) + monkeypatch.setattr(FakeScene, "position_limits_lower", [-2.0, -1.0, -3.0]) + monkeypatch.setattr(FakeScene, "position_limits_upper", [2.0, 1.0, 3.0]) + fk_frames: list[str] = [] + + def fake_fk( + self: FakeScene, q: np.ndarray, frame_name: str, base_frame: str = "" + ) -> np.ndarray: + _ = (self, base_frame) + fk_frames.append(frame_name) + mat = np.eye(4) + mat[0, 3] = float(np.sum(q)) + return mat + + def fake_jacobian( + self: FakeScene, q: np.ndarray, frame_name: str, local: bool = True + ) -> np.ndarray: + _ = (self, q) + assert frame_name == "wrist_tip" + assert local is True + return np.arange(18, dtype=np.float64).reshape(6, 3) + + monkeypatch.setattr(FakeScene, "forwardKinematics", fake_fk) + monkeypatch.setattr(FakeScene, "computeFrameJacobian", fake_jacobian) + world, robot_id = _make_world(fake_roboplan, config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, + robot_id, + JointState({"name": ["joint1", "joint2", "joint3"], "position": [1.0, 2.0, 3.0]}), + ) + + pose = world.get_group_ee_pose(ctx, "arm/wrist") + jacobian = world.get_group_jacobian(ctx, "arm/wrist") + + assert fk_frames == ["wrist_tip"] + assert pose.position.x == pytest.approx(6.0) + np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 1]]) + + def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world, robot_id = _make_world(fake_roboplan, robot_config) world.finalize() diff --git a/openspec/changes/planning-world-monitor-groups/tasks.md b/openspec/changes/planning-world-monitor-groups/tasks.md index acb647d94a..8e31566a05 100644 --- a/openspec/changes/planning-world-monitor-groups/tasks.md +++ b/openspec/changes/planning-world-monitor-groups/tasks.md @@ -1,18 +1,19 @@ ## 1. Extract backend group support -- [ ] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. -- [ ] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. -- [ ] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. -- [ ] 1.4 Extract world monitor and robot state monitor group-aware query changes. +- [x] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. +- [x] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. +- [x] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. +- [x] 1.4 Extract world monitor and robot state monitor group-aware query changes. ## 2. Tests -- [ ] 2.1 Bring over Drake group world tests. -- [ ] 2.2 Bring over RoboPlan world tests, including planning split file if needed. -- [ ] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. +- [x] 2.1 Bring over Drake group world tests. +- [x] 2.2 Bring over RoboPlan world tests, including planning split file if needed. +- [x] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. ## 3. Validation -- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. -- [ ] 3.2 Run targeted mypy on changed backend/monitor files. -- [ ] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. +- [x] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. + - Note: `dimos/manipulation/test_roboplan_world_planning.py` does not exist in this branch; ran the same targeted command without that optional split file. +- [x] 3.2 Run targeted mypy on changed backend/monitor files. +- [x] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. From e59b9ef66eb9f07d4a8bd77b3927d95dd457d11a Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 22:09:34 -0700 Subject: [PATCH 05/29] refactor: centralize planning group registry --- dimos/manipulation/planning/groups/utils.py | 149 +++++++++++++++++- .../planning/monitor/test_world_monitor.py | 58 ++++++- .../planning/monitor/world_monitor.py | 93 +++-------- .../planning/world/drake_world.py | 39 ++--- .../planning/world/roboplan_world.py | 32 ++-- 5 files changed, 264 insertions(+), 107 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index d128d18e3a..b642251471 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -20,12 +20,17 @@ assert_global_joint_names, assert_local_joint_names, is_global_joint_name, + make_global_joint_names, + make_planning_group_id, + parse_planning_group_id, ) -from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( GlobalJointName, LocalModelJointName, PlanningGroupID, + RobotName, ) from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -55,6 +60,144 @@ def matching_global_joint_name( return None +def planning_group_from_configs( + group_id: PlanningGroupID, configs: Sequence[RobotModelConfig] +) -> PlanningGroup: + """Resolve a public planning-group ID from stored robot configs.""" + robot_name, group_name = parse_planning_group_id(group_id) + config = _config_for_robot_name(robot_name, configs) + definition = _definition_for_group_name(group_name, config.planning_groups, group_id) + return planning_group_from_definition(config, definition) + + +def planning_group_from_definition( + config: RobotModelConfig, definition: PlanningGroupDefinition +) -> PlanningGroup: + """Build the public planning-group model for one config definition.""" + return PlanningGroup( + id=make_planning_group_id(config.name, definition.name), + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + + +def validate_planning_group_config(config: RobotModelConfig) -> None: + """Validate config-derived planning groups without constructing a registry.""" + seen_group_ids: set[PlanningGroupID] = set() + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names or group_id in seen_group_ids: + raise ValueError(f"Planning group '{group_id}' is already registered") + planning_group_from_definition(config, definition) + seen_group_names.add(definition.name) + seen_group_ids.add(group_id) + + +def primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: + """Return the unique pose-targetable group ID for robot-scoped wrappers.""" + pose_groups = [group for group in config.planning_groups if group.has_pose_target] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError( + f"Robot '{config.name}' has {len(pose_groups)} pose-targetable planning groups; " + "use an explicit planning group ID" + ) + return make_planning_group_id(config.name, pose_groups[0].name) + + +def full_robot_joint_state_from_input( + config: RobotModelConfig, joint_state: JointState +) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + +def joint_state_for_group_query( + config: RobotModelConfig, + group: PlanningGroup, + current_state: JointState, + joint_state: JointState | None, +) -> JointState: + """Normalize full-robot or group-scoped input to full robot-local state.""" + if joint_state is None: + return current_state + full_state = full_robot_joint_state_from_input(config, joint_state) + if full_state is not None: + return full_state + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names + ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + + +def _config_for_robot_name( + robot_name: RobotName, configs: Sequence[RobotModelConfig] +) -> RobotModelConfig: + for config in configs: + if config.name == robot_name: + return config + raise KeyError(f"No robot registered for planning group robot '{robot_name}'") + + +def _definition_for_group_name( + group_name: str, + definitions: Sequence[PlanningGroupDefinition], + group_id: PlanningGroupID, +) -> PlanningGroupDefinition: + for definition in definitions: + if definition.name == group_name: + return definition + raise KeyError(f"Unknown planning group ID: {group_id}") + + def filter_joint_state_to_selected_joints( joint_state: JointState, global_joint_names: Sequence[GlobalJointName], @@ -104,7 +247,7 @@ def joint_target_to_global_names( f"Target for '{group.id}' has {len(target.position)} positions, " f"expected {len(group.joint_names)}" ) - return JointState(name=list(group.joint_names), position=list(target.position)) + return JointState({"name": list(group.joint_names), "position": list(target.position)}) if len(target.name) != len(target.position): raise ValueError( @@ -140,4 +283,4 @@ def joint_target_to_global_names( extra = set(target_names) - set(expected_names) if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") - return JointState(name=list(group.joint_names), position=global_positions) + return JointState({"name": list(group.joint_names), "position": global_positions}) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index aed6d6a9a5..b7d9016922 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -288,18 +288,74 @@ def test_world_monitor_exposes_planning_groups_and_duplicate_names_do_not_mutate assert [call[0] for call in fake_world.calls].count("add_robot") == 1 +def test_world_monitor_invalid_duplicate_group_config_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1",), base_link="base", tip_link="ee" + ), + PlanningGroupDefinition( + name="manipulator", joint_names=("j2",), base_link="base", tip_link="ee" + ), + ] + ) + + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + +def test_world_monitor_invalid_group_joint_name_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j1", "bad/joint"), + base_link="base", + tip_link="ee", + ) + ] + ) + + with pytest.raises(ValueError, match="Invalid local joint name"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + def test_current_group_joint_state_uses_public_names_in_group_order() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] robot_id = monitor.add_robot(_three_joint_reordered_group_config()) monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] - state = monitor.get_current_group_joint_state("arm/manipulator") + state = monitor.current_group_joint_state("arm/manipulator") assert state.name == ["arm/j2", "arm/j1"] assert state.position == [0.2, 0.1] +def test_group_kinematics_with_full_state_does_not_require_current_state() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_three_joint_reordered_group_config()) + + pose = monitor.get_group_ee_pose( + "arm/manipulator", + JointState({"name": ["j1", "j2", "j3"], "position": [0.1, 0.2, 0.3]}), + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 3cd54d05c4..13ff2d2fc3 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -22,7 +22,12 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT 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.groups.utils import ( + filter_joint_state_to_selected_joints, + full_robot_joint_state_from_input, + joint_state_for_group_query, + validate_planning_group_config, +) from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.spec.models import PlanningSceneInfo @@ -80,8 +85,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: with self._lock: if config.name in self._robot_ids_by_name: raise ValueError(f"Robot name '{config.name}' is already registered") - new_groups = PlanningGroupRegistry() - new_groups.add_robot(config) + validate_planning_group_config(config) robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names self._robot_configs[robot_id] = config @@ -301,9 +305,11 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: if positions is not None: joint_names = self._robot_joints.get(robot_id, []) return JointState( - name=joint_names, - position=positions.tolist(), - velocity=velocities.tolist() if velocities is not None else [], + { + "name": joint_names, + "position": positions.tolist(), + "velocity": velocities.tolist() if velocities is not None else [], + } ) # Fall back to world's live context @@ -328,17 +334,12 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: positions.append(float(pos_by_name[local_name])) return JointState({"name": names, "position": positions}) - def get_current_group_joint_state( + def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: """Return current joint state scoped and ordered for one planning group.""" group = self._planning_groups.get(group_id) - robot_id = self._robot_ids_by_name[group.robot_name] - if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): - raise ValueError(f"Current state for robot '{group.robot_name}' is stale") - state = self.get_current_joint_state(robot_id) - if state is None: - raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + state = self._current_robot_joint_state_for_group(group_id, max_age) return filter_joint_state_to_selected_joints( state, group.joint_names, group.local_joint_names ) @@ -356,69 +357,21 @@ def _current_robot_joint_state_for_group( raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") return state - def _full_robot_joint_state_from_input( - self, robot_id: WorldRobotID, joint_state: JointState - ) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - config = self._robot_configs[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } - ) - def _joint_state_for_group_query( self, group_id: PlanningGroupID, joint_state: JointState | None ) -> JointState: """Normalize full-robot or group-scoped input to full robot-local state.""" group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] - if joint_state is None: - return self._current_robot_joint_state_for_group(group_id) - full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) - if full_state is not None: - return full_state + if joint_state is not None: + full_state = full_robot_joint_state_from_input( + self._robot_configs[robot_id], joint_state + ) + if full_state is not None: + return full_state current_state = self._current_robot_joint_state_for_group(group_id) - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } + return joint_state_for_group_query( + self._robot_configs[robot_id], group, current_state, joint_state ) def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: @@ -427,7 +380,7 @@ def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: velocities = self._state_monitors[robot_id].get_current_velocities() if velocities is not None: joint_names = self._robot_joints.get(robot_id, []) - return JointState(name=joint_names, velocity=velocities.tolist()) + return JointState({"name": joint_names, "velocity": velocities.tolist()}) return None def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 12067f647c..b1aa4da024 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -26,7 +26,11 @@ import numpy as np from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + planning_group_from_configs, + primary_pose_group_id_for_config, + validate_planning_group_config, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -187,7 +191,6 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: # Tracking data self._robots: dict[WorldRobotID, _RobotData] = {} - self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, _ObstacleData] = {} self._robot_counter = 0 self._obstacle_counter = 0 @@ -211,10 +214,9 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: - if self._planning_groups.groups_for_robot(config.name): + if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") - new_groups = PlanningGroupRegistry() - new_groups.add_robot(config) + validate_planning_group_config(config) self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -226,11 +228,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: ee_link = config.base_link try: - primary_group_id = new_groups.primary_pose_group_id_for_robot(config.name) + primary_group_id = primary_pose_group_id_for_config(config) except ValueError: primary_group_id = None if primary_group_id is not None: - primary_group = new_groups.get(primary_group_id) + primary_group = planning_group_from_configs(primary_group_id, [config]) if primary_group.tip_link is not None: ee_link = primary_group.tip_link ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() @@ -251,8 +253,6 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: base_frame=base_frame, preview_model_instance=preview_model_instance, ) - self._planning_groups.add_robot(config) - logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id @@ -334,11 +334,6 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config - @property - def planning_groups(self) -> PlanningGroupRegistry: - """Registered planning groups for this world.""" - return self._planning_groups - def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -866,7 +861,9 @@ def _joint_state_to_q( return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] @@ -973,7 +970,7 @@ def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + group_id = primary_pose_group_id_for_config(robot_data.config) if group_id is None: raise ValueError( f"Robot '{robot_data.config.name}' has no pose-targetable planning group" @@ -985,7 +982,9 @@ def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStam if not self._finalized: raise RuntimeError("World must be finalized first") - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_data = self._robots[self._robot_id_for_group(group_id)] @@ -1035,7 +1034,7 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + group_id = primary_pose_group_id_for_config(robot_data.config) if group_id is None: raise ValueError( f"Robot '{robot_data.config.name}' has no pose-targetable planning group" @@ -1047,7 +1046,9 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray if not self._finalized: raise RuntimeError("World must be finalized first") - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_data = self._robots[self._robot_id_for_group(group_id)] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index e224ff583a..11a19c883c 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -42,7 +42,11 @@ ) from exc from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + planning_group_from_configs, + primary_pose_group_id_for_config, + validate_planning_group_config, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -95,7 +99,6 @@ def __init__(self, enable_viz: bool = False, **_: object) -> None: logger.warning("RoboPlanWorld does not currently provide manipulation visualization") self._robots: dict[WorldRobotID, _RoboPlanRobotData] = {} - self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, Obstacle] = {} self._robot_counter = 0 self._finalized = False @@ -112,8 +115,9 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise ValueError("RoboPlanWorld currently supports one robot per Scene") if not Path(config.model_path).exists(): raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") - if self._planning_groups.groups_for_robot(config.name): + if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") + validate_planning_group_config(config) self._validate_robot_config(config) self._robot_counter += 1 @@ -129,7 +133,6 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._live_context.q_by_robot[robot_id] = np.zeros( len(config.joint_names), dtype=np.float64 ) - self._planning_groups.add_robot(config) logger.info(f"Added RoboPlan robot '{robot_id}' ({config.name})") return robot_id @@ -141,11 +144,6 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration by ID.""" return self._get_robot(robot_id).config - @property - def planning_groups(self) -> PlanningGroupRegistry: - """Registered planning groups for this world.""" - return self._planning_groups - def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -288,14 +286,16 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + group_id = primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_ee_pose(ctx, group_id) def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: """Get planning-group tip pose if RoboPlan exposes FK.""" - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") mat = self.get_link_pose(ctx, self._robot_id_for_group(group_id), group.tip_link) @@ -325,7 +325,7 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + group_id = primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_jacobian(ctx, group_id) @@ -334,7 +334,9 @@ def get_group_jacobian( self, ctx: RoboPlanContext, group_id: PlanningGroupID ) -> NDArray[np.float64]: """Get planning-group Jacobian projected to group-local joint order.""" - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_id = self._robot_id_for_group(group_id) @@ -563,7 +565,9 @@ def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: return self._robots[robot_id] def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] From db1b72902b1a9a70a24abdee516fcd202672d32e Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 22:45:00 -0700 Subject: [PATCH 06/29] refactor: align planning group utilities --- .../planning/groups/test_planning_groups.py | 18 -- dimos/manipulation/planning/groups/utils.py | 158 +----------------- .../planning/monitor/world_monitor.py | 82 +++++++-- .../planning/world/drake_world.py | 75 ++++++--- .../planning/world/roboplan_world.py | 78 +++++++-- .../world/test_drake_world_planning_groups.py | 2 +- 6 files changed, 192 insertions(+), 221 deletions(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index ae5776fd74..e716366202 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -20,7 +20,6 @@ import pytest -from dimos.manipulation.planning.groups import utils as planning_group_utils from dimos.manipulation.planning.groups.discovery import ( FALLBACK_PLANNING_GROUP_NAME, PlanningGroupDiscoveryError, @@ -474,23 +473,6 @@ def test_matching_global_joint_name_requires_unique_suffix_match() -> None: assert matching_global_joint_name({"left/j1": 1.0}, "j2") is None -def test_matching_global_joint_name_warns_when_suffix_match_is_not_unique( - monkeypatch: pytest.MonkeyPatch, -) -> None: - messages: list[str] = [] - - def capture_warning(message: str) -> None: - messages.append(message) - - monkeypatch.setattr(planning_group_utils.logger, "warning", capture_warning) - - assert matching_global_joint_name({"left/j1": 1.0, "right/j1": 2.0}, "j1") is None - assert messages == [ - "Expected exactly one global joint ending with local joint name 'j1', " - "found 2 matches: ['left/j1', 'right/j1']" - ] - - def test_planning_group_id_from_selector_accepts_id_or_group() -> None: group = _make_group() diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index b642251471..e55f997a33 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -20,22 +20,14 @@ assert_global_joint_names, assert_local_joint_names, is_global_joint_name, - make_global_joint_names, - make_planning_group_id, - parse_planning_group_id, ) -from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition -from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import ( GlobalJointName, LocalModelJointName, PlanningGroupID, - RobotName, ) from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() def planning_group_id_from_selector(selector: PlanningGroupID | PlanningGroup) -> PlanningGroupID: @@ -53,151 +45,9 @@ def matching_global_joint_name( matches = [name for name in positions_by_name if name.endswith(suffix)] if len(matches) == 1: return matches[0] - logger.warning( - f"Expected exactly one global joint ending with local joint name " - f"'{local_joint_name}', found {len(matches)} matches: {matches}" - ) return None -def planning_group_from_configs( - group_id: PlanningGroupID, configs: Sequence[RobotModelConfig] -) -> PlanningGroup: - """Resolve a public planning-group ID from stored robot configs.""" - robot_name, group_name = parse_planning_group_id(group_id) - config = _config_for_robot_name(robot_name, configs) - definition = _definition_for_group_name(group_name, config.planning_groups, group_id) - return planning_group_from_definition(config, definition) - - -def planning_group_from_definition( - config: RobotModelConfig, definition: PlanningGroupDefinition -) -> PlanningGroup: - """Build the public planning-group model for one config definition.""" - return PlanningGroup( - id=make_planning_group_id(config.name, definition.name), - robot_name=config.name, - group_name=definition.name, - joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), - local_joint_names=definition.joint_names, - base_link=definition.base_link, - tip_link=definition.tip_link, - source=definition.source, - ) - - -def validate_planning_group_config(config: RobotModelConfig) -> None: - """Validate config-derived planning groups without constructing a registry.""" - seen_group_ids: set[PlanningGroupID] = set() - seen_group_names: set[str] = set() - for definition in config.planning_groups: - group_id = make_planning_group_id(config.name, definition.name) - if definition.name in seen_group_names or group_id in seen_group_ids: - raise ValueError(f"Planning group '{group_id}' is already registered") - planning_group_from_definition(config, definition) - seen_group_names.add(definition.name) - seen_group_ids.add(group_id) - - -def primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: - """Return the unique pose-targetable group ID for robot-scoped wrappers.""" - pose_groups = [group for group in config.planning_groups if group.has_pose_target] - if not pose_groups: - return None - if len(pose_groups) > 1: - raise ValueError( - f"Robot '{config.name}' has {len(pose_groups)} pose-targetable planning groups; " - "use an explicit planning group ID" - ) - return make_planning_group_id(config.name, pose_groups[0].name) - - -def full_robot_joint_state_from_input( - config: RobotModelConfig, joint_state: JointState -) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } - ) - - -def joint_state_for_group_query( - config: RobotModelConfig, - group: PlanningGroup, - current_state: JointState, - joint_state: JointState | None, -) -> JointState: - """Normalize full-robot or group-scoped input to full robot-local state.""" - if joint_state is None: - return current_state - full_state = full_robot_joint_state_from_input(config, joint_state) - if full_state is not None: - return full_state - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } - ) - - -def _config_for_robot_name( - robot_name: RobotName, configs: Sequence[RobotModelConfig] -) -> RobotModelConfig: - for config in configs: - if config.name == robot_name: - return config - raise KeyError(f"No robot registered for planning group robot '{robot_name}'") - - -def _definition_for_group_name( - group_name: str, - definitions: Sequence[PlanningGroupDefinition], - group_id: PlanningGroupID, -) -> PlanningGroupDefinition: - for definition in definitions: - if definition.name == group_name: - return definition - raise KeyError(f"Unknown planning group ID: {group_id}") - - def filter_joint_state_to_selected_joints( joint_state: JointState, global_joint_names: Sequence[GlobalJointName], @@ -226,7 +76,7 @@ def filter_joint_state_to_selected_joints( missing.append(global_name) if missing: - raise ValueError(f"Joint state is missing selected joints: {missing}") + raise ValueError(f"IK result is missing selected joints: {missing}") return JointState({"name": list(global_joint_names), "position": selected_positions}) @@ -247,7 +97,7 @@ def joint_target_to_global_names( f"Target for '{group.id}' has {len(target.position)} positions, " f"expected {len(group.joint_names)}" ) - return JointState({"name": list(group.joint_names), "position": list(target.position)}) + return JointState(name=list(group.joint_names), position=list(target.position)) if len(target.name) != len(target.position): raise ValueError( @@ -283,4 +133,4 @@ def joint_target_to_global_names( extra = set(target_names) - set(expected_names) if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") - return JointState({"name": list(group.joint_names), "position": global_positions}) + return JointState(name=list(group.joint_names), position=global_positions) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 13ff2d2fc3..cf38697678 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -21,13 +21,12 @@ from typing import TYPE_CHECKING, Any from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry -from dimos.manipulation.planning.groups.utils import ( - filter_joint_state_to_selected_joints, - full_robot_joint_state_from_input, - joint_state_for_group_query, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, ) +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.spec.models import PlanningSceneInfo @@ -85,7 +84,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: with self._lock: if config.name in self._robot_ids_by_name: raise ValueError(f"Robot name '{config.name}' is already registered") - validate_planning_group_config(config) + self._validate_planning_group_config(config) robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names self._robot_configs[robot_id] = config @@ -364,15 +363,74 @@ def _joint_state_for_group_query( group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] if joint_state is not None: - full_state = full_robot_joint_state_from_input( - self._robot_configs[robot_id], joint_state - ) + full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) if full_state is not None: return full_state current_state = self._current_robot_joint_state_for_group(group_id) - return joint_state_for_group_query( - self._robot_configs[robot_id], group, current_state, joint_state + if joint_state is None: + return current_state + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + + def _full_robot_joint_state_from_input( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + config = self._robot_configs[robot_id] + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + """Validate planning groups before mutating world/backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index b1aa4da024..b84c735556 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -25,12 +25,12 @@ import numpy as np -from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.utils import ( - planning_group_from_configs, - primary_pose_group_id_for_config, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + make_global_joint_names, + make_planning_group_id, ) +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -216,7 +216,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: with self._lock: if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") - validate_planning_group_config(config) + self._validate_planning_group_config(config) self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -228,11 +228,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: ee_link = config.base_link try: - primary_group_id = primary_pose_group_id_for_config(config) + primary_group_id = self._primary_pose_group_id_for_config(config) except ValueError: primary_group_id = None if primary_group_id is not None: - primary_group = planning_group_from_configs(primary_group_id, [config]) + primary_group = self._planning_group_from_config(config, primary_group_id) if primary_group.tip_link is not None: ee_link = primary_group.tip_link ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() @@ -334,6 +334,49 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + seen_group_names: set[str] = set() + for definition in config.planning_groups: + make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{definition.name}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + def _planning_group_from_config( + self, config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + joint_names = tuple(make_global_joint_names(config.name, definition.joint_names)) + return PlanningGroup( + group_id, + config.name, + definition.name, + joint_names, + definition.joint_names, + definition.base_link, + definition.tip_link, + definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot_data in self._robots.values(): + try: + return self._planning_group_from_config(robot_data.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + pose_groups = [group for group in config.planning_groups if group.has_pose_target] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError(f"Robot '{config.name}' has multiple pose groups") + return make_planning_group_id(config.name, pose_groups[0].name) + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -861,9 +904,7 @@ def _joint_state_to_q( return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] @@ -970,7 +1011,7 @@ def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = primary_pose_group_id_for_config(robot_data.config) + group_id = self._primary_pose_group_id_for_config(robot_data.config) if group_id is None: raise ValueError( f"Robot '{robot_data.config.name}' has no pose-targetable planning group" @@ -982,9 +1023,7 @@ def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStam if not self._finalized: raise RuntimeError("World must be finalized first") - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_data = self._robots[self._robot_id_for_group(group_id)] @@ -1034,7 +1073,7 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = primary_pose_group_id_for_config(robot_data.config) + group_id = self._primary_pose_group_id_for_config(robot_data.config) if group_id is None: raise ValueError( f"Robot '{robot_data.config.name}' has no pose-targetable planning group" @@ -1046,9 +1085,7 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray if not self._finalized: raise RuntimeError("World must be finalized first") - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_data = self._robots[self._robot_id_for_group(group_id)] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 11a19c883c..0a3c69d30c 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -41,12 +41,12 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc -from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.utils import ( - planning_group_from_configs, - primary_pose_group_id_for_config, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + make_global_joint_names, + make_planning_group_id, ) +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -117,7 +117,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") - validate_planning_group_config(config) + self._validate_planning_group_config(config) self._validate_robot_config(config) self._robot_counter += 1 @@ -286,16 +286,14 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - group_id = primary_pose_group_id_for_config(robot.config) + group_id = self._primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_ee_pose(ctx, group_id) def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: """Get planning-group tip pose if RoboPlan exposes FK.""" - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") mat = self.get_link_pose(ctx, self._robot_id_for_group(group_id), group.tip_link) @@ -325,7 +323,7 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) - group_id = primary_pose_group_id_for_config(robot.config) + group_id = self._primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_jacobian(ctx, group_id) @@ -334,9 +332,7 @@ def get_group_jacobian( self, ctx: RoboPlanContext, group_id: PlanningGroupID ) -> NDArray[np.float64]: """Get planning-group Jacobian projected to group-local joint order.""" - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) if group.tip_link is None: raise ValueError(f"Planning group '{group_id}' has no tip link") robot_id = self._robot_id_for_group(group_id) @@ -559,15 +555,63 @@ def _query_scene_joint_order( return None return list(group_info.joint_names) + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + """Validate planning groups before mutating backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + def _planning_group_from_config( + self, config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + return PlanningGroup( + id=group_id, + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot in self._robots.values(): + try: + return self._planning_group_from_config(robot.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + pose_group_ids = [ + make_planning_group_id(config.name, group.name) + for group in config.planning_groups + if group.has_pose_target + ] + if not pose_group_ids: + return None + if len(pose_group_ids) > 1: + raise ValueError( + f"Robot '{config.name}' has {len(pose_group_ids)} pose-targetable " + "planning groups; use an explicit planning group ID" + ) + return pose_group_ids[0] + def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id] def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index 841ed27f99..1331e447b7 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -138,7 +138,7 @@ def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_pa ) ) ambiguous.finalize() - with pytest.raises(ValueError, match="multiple pose-targetable"): + with pytest.raises(ValueError, match="multiple pose"): ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) From f9b269c848c53c8519a15a93b8e9d3f1c757ca5c Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 10:53:25 -0700 Subject: [PATCH 07/29] chore: revert kwarg call --- .../planning/monitor/world_monitor.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index cf38697678..5d5f1f4707 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -304,11 +304,9 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: if positions is not None: joint_names = self._robot_joints.get(robot_id, []) return JointState( - { - "name": joint_names, - "position": positions.tolist(), - "velocity": velocities.tolist() if velocities is not None else [], - } + name=joint_names, + position=positions.tolist(), + velocity=velocities.tolist() if velocities is not None else [], ) # Fall back to world's live context @@ -331,7 +329,7 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: if local_name in pos_by_name: names.append(f"{robot_name}/{local_name}") positions.append(float(pos_by_name[local_name])) - return JointState({"name": names, "position": positions}) + return JointState(name=names, position=positions) def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 @@ -378,10 +376,8 @@ def _joint_state_for_group_query( raise ValueError(f"Current state is missing group joint '{local_name}'") positions_by_name[local_name] = float(position) return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } + name=list(current_state.name), + position=[positions_by_name[name] for name in current_state.name], ) def _full_robot_joint_state_from_input( @@ -392,9 +388,7 @@ def _full_robot_joint_state_from_input( if not joint_state.name: if len(joint_state.position) != len(config.joint_names): return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) + return JointState(name=list(config.joint_names), position=list(joint_state.position)) if len(joint_state.name) != len(joint_state.position): raise ValueError("JointState name and position lengths must match") resolved_positions: dict[str, float] = {} @@ -416,10 +410,8 @@ def _full_robot_joint_state_from_input( if set(resolved_positions) != set(config.joint_names): return None return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } + name=list(config.joint_names), + position=[resolved_positions[name] for name in config.joint_names], ) def _validate_planning_group_config(self, config: RobotModelConfig) -> None: @@ -438,7 +430,7 @@ def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: velocities = self._state_monitors[robot_id].get_current_velocities() if velocities is not None: joint_names = self._robot_joints.get(robot_id, []) - return JointState({"name": joint_names, "velocity": velocities.tolist()}) + return JointState(name=joint_names, velocity=velocities.tolist()) return None def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: From 8d0527112967c833addaf771737f12d0563744ca Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 11:34:25 -0700 Subject: [PATCH 08/29] refactor: simplify group state queries --- .../planning/monitor/test_world_monitor.py | 19 ++--- .../planning/monitor/world_monitor.py | 84 +++---------------- 2 files changed, 22 insertions(+), 81 deletions(-) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index b7d9016922..84a17160d1 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -347,7 +347,7 @@ def test_group_kinematics_with_full_state_does_not_require_current_state() -> No pose = monitor.get_group_ee_pose( "arm/manipulator", - JointState({"name": ["j1", "j2", "j3"], "position": [0.1, 0.2, 0.3]}), + JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3]), ) set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] @@ -356,17 +356,18 @@ def test_group_kinematics_with_full_state_does_not_require_current_state() -> No assert pose.position.x == 1 -def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: +def test_group_kinematics_route_full_state_to_backend() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] - robot_id = monitor.add_robot(_three_joint_reordered_group_config()) - monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + monitor.add_robot(_three_joint_reordered_group_config()) pose = monitor.get_group_ee_pose( - "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.8, 0.9]}) + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.9, 0.8, 0.3]), ) jacobian = monitor.get_group_jacobian( - "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.3, 0.4]}) + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.4, 0.3, 0.3]), ) set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] @@ -389,7 +390,7 @@ def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: ) ) with pytest.raises(ValueError, match="no pose-targetable"): - monitor.get_ee_pose(no_pose_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]})) + monitor.get_ee_pose(no_pose_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) fake_world2 = FakeWorld() monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] @@ -406,6 +407,4 @@ def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: ) ) with pytest.raises(ValueError, match="pose-targetable planning groups"): - monitor2.get_jacobian( - ambiguous_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]}) - ) + monitor2.get_jacobian(ambiguous_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 5d5f1f4707..9a0241d6a2 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -335,16 +335,12 @@ def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: """Return current joint state scoped and ordered for one planning group.""" - group = self._planning_groups.get(group_id) - state = self._current_robot_joint_state_for_group(group_id, max_age) - return filter_joint_state_to_selected_joints( - state, group.joint_names, group.local_joint_names - ) + return self._current_robot_joint_state_for_group(group_id, max_age) def _current_robot_joint_state_for_group( self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: - """Return current full robot state for a planning group.""" + """Return current joint state scoped and ordered for one planning group.""" group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): @@ -352,66 +348,8 @@ def _current_robot_joint_state_for_group( state = self.get_current_joint_state(robot_id) if state is None: raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") - return state - - def _joint_state_for_group_query( - self, group_id: PlanningGroupID, joint_state: JointState | None - ) -> JointState: - """Normalize full-robot or group-scoped input to full robot-local state.""" - group = self._planning_groups.get(group_id) - robot_id = self._robot_ids_by_name[group.robot_name] - if joint_state is not None: - full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) - if full_state is not None: - return full_state - current_state = self._current_robot_joint_state_for_group(group_id) - if joint_state is None: - return current_state - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - name=list(current_state.name), - position=[positions_by_name[name] for name in current_state.name], - ) - - def _full_robot_joint_state_from_input( - self, robot_id: WorldRobotID, joint_state: JointState - ) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - config = self._robot_configs[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState(name=list(config.joint_names), position=list(joint_state.position)) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - name=list(config.joint_names), - position=[resolved_positions[name] for name in config.joint_names], + return filter_joint_state_to_selected_joints( + state, group.joint_names, group.local_joint_names ) def _validate_planning_group_config(self, config: RobotModelConfig) -> None: @@ -506,12 +444,17 @@ def get_ee_pose( def get_group_ee_pose( self, group_id: PlanningGroupID, joint_state: JointState | None = None ) -> PoseStamped: - """Get planning-group tip pose. Uses current group state if joint_state is None.""" + """Get planning-group tip pose. Uses current robot state if joint_state is None.""" group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - normalized_state = self._joint_state_for_group_query(group_id, joint_state) - self._world.set_joint_state(ctx, robot_id, normalized_state) + if joint_state is None: + if robot_id in self._state_monitors and self.is_state_stale(robot_id): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + joint_state = self.get_current_joint_state(robot_id) + if joint_state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_group_ee_pose(ctx, group_id) @@ -562,8 +505,7 @@ def get_group_jacobian( group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - normalized_state = self._joint_state_for_group_query(group_id, joint_state) - self._world.set_joint_state(ctx, robot_id, normalized_state) + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_group_jacobian(ctx, group_id) # Lifecycle From 340aef7b906dbb2fc66f2a8a271027fe6610cefe Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:04:50 -0700 Subject: [PATCH 09/29] refactor: simplify current group state --- .../manipulation/planning/monitor/world_monitor.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 9a0241d6a2..9b54332de8 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -324,21 +324,13 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: state = self.get_current_joint_state(robot_id) if state is None: continue - pos_by_name = dict(zip(state.name, state.position, strict=True)) - for local_name in self._robot_joints.get(robot_id, []): - if local_name in pos_by_name: - names.append(f"{robot_name}/{local_name}") - positions.append(float(pos_by_name[local_name])) + for name, position in zip(state.name, state.position, strict=True): + names.append(f"{robot_name}/{name}") + positions.append(float(position)) return JointState(name=names, position=positions) def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 - ) -> JointState: - """Return current joint state scoped and ordered for one planning group.""" - return self._current_robot_joint_state_for_group(group_id, max_age) - - def _current_robot_joint_state_for_group( - self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: """Return current joint state scoped and ordered for one planning group.""" group = self._planning_groups.get(group_id) From a33f8ec0e147e45c571570915ba127e9ccbc70b9 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:26:16 -0700 Subject: [PATCH 10/29] refactor: share joint state ordering --- dimos/manipulation/planning/groups/utils.py | 47 +++++++++++++++++++ .../planning/world/drake_world.py | 37 ++++----------- .../planning/world/roboplan_world.py | 35 +++----------- 3 files changed, 62 insertions(+), 57 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index e55f997a33..ae3e865c8f 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -16,6 +16,9 @@ from collections.abc import Mapping, Sequence +import numpy as np +from numpy.typing import NDArray + from dimos.manipulation.planning.groups.identifiers import ( assert_global_joint_names, assert_local_joint_names, @@ -134,3 +137,47 @@ def joint_target_to_global_names( if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") return JointState(name=list(group.joint_names), position=global_positions) + + +def joint_state_to_ordered_positions( + joint_state: JointState, + *, + robot_name: str, + joint_names: Sequence[str], + joint_name_mapping: Mapping[str, str], + context: str, +) -> NDArray[np.float64]: + """Convert a JointState to an array ordered by local robot joint names.""" + if not joint_state.name: + if len(joint_state.position) != len(joint_names): + raise ValueError("JointState position length must match configured joint count") + return np.asarray(joint_state.position, dtype=np.float64) + + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + + joint_name_set = set(joint_names) + name_to_pos: dict[str, float] = {} + prefix = f"{robot_name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in joint_name_set: + resolved_name = name + elif name in joint_name_mapping: + resolved_name = joint_name_mapping[name] + elif is_global_joint_name(name): + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in joint_name_set: + raise ValueError(f"Unknown global joint name for {context}: {name}") + else: + resolved_name = joint_name_mapping.get(name, name) + + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) + + missing = [name for name in joint_names if name not in name_to_pos] + if missing: + raise ValueError(f"JointState missing joints for {context}: {missing}") + return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index b84c735556..e5cb3cb72a 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -26,11 +26,11 @@ import numpy as np from dimos.manipulation.planning.groups.identifiers import ( - is_global_joint_name, make_global_joint_names, make_planning_group_id, ) from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -874,34 +874,13 @@ def _joint_state_to_q( if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(robot_data.config.joint_names): - raise ValueError("JointState position length must match configured joint count") - return np.asarray(joint_state.position, dtype=np.float64) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - name_to_pos: dict[str, float] = {} - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in robot_data.config.joint_names: - resolved_name = name - elif name in robot_data.config.joint_name_mapping: - resolved_name = robot_data.config.joint_name_mapping[name] - elif is_global_joint_name(name): - prefix = f"{robot_data.config.name}/" - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] - if resolved_name not in robot_data.config.joint_names: - raise ValueError(f"Unknown global joint name for DrakeWorld: {name}") - else: - resolved_name = robot_data.config.get_urdf_joint_name(name) - if resolved_name in name_to_pos: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - name_to_pos[resolved_name] = float(position) - missing = [name for name in robot_data.config.joint_names if name not in name_to_pos] - if missing: - raise ValueError(f"JointState missing joints for DrakeWorld: {missing}") - return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) + return joint_state_to_ordered_positions( + joint_state, + robot_name=robot_data.config.name, + joint_names=robot_data.config.joint_names, + joint_name_mapping=robot_data.config.joint_name_mapping, + context="DrakeWorld", + ) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: group = self._planning_group_from_id(group_id) diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 0a3c69d30c..5038fdeebe 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -42,11 +42,11 @@ ) from exc from dimos.manipulation.planning.groups.identifiers import ( - is_global_joint_name, make_global_joint_names, make_planning_group_id, ) from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -623,33 +623,12 @@ def _joint_state_to_q( self, robot_id: WorldRobotID, joint_state: JointState ) -> NDArray[np.float64]: robot = self._get_robot(robot_id) - if len(joint_state.position) != len(robot.config.joint_names): - raise ValueError("JointState position length must match configured joint count") - if not joint_state.name: - return np.asarray(joint_state.position, dtype=np.float64) - name_to_pos: dict[str, float] = {} - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in robot.config.joint_names: - resolved_name = name - elif name in robot.config.joint_name_mapping: - resolved_name = robot.config.joint_name_mapping[name] - elif is_global_joint_name(name): - prefix = f"{robot.config.name}/" - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] - if resolved_name not in robot.config.joint_names: - raise ValueError(f"Unknown global joint name for RoboPlanWorld: {name}") - else: - resolved_name = robot.config.get_urdf_joint_name(name) - if resolved_name in name_to_pos: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - name_to_pos[resolved_name] = float(position) - missing = [name for name in robot.config.joint_names if name not in name_to_pos] - if missing: - raise ValueError(f"JointState missing joints for RoboPlanWorld: {missing}") - return np.asarray( - [name_to_pos[name] for name in robot.config.joint_names], dtype=np.float64 + return joint_state_to_ordered_positions( + joint_state, + robot_name=robot.config.name, + joint_names=robot.config.joint_names, + joint_name_mapping=robot.config.joint_name_mapping, + context="RoboPlanWorld", ) def _require_finalized(self) -> None: From df05bcbfa13149e99841804ed0989162cfbfabe2 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:32:30 -0700 Subject: [PATCH 11/29] refactor: simplify joint state ordering util --- dimos/manipulation/planning/groups/utils.py | 11 +++-------- dimos/manipulation/planning/world/drake_world.py | 2 -- dimos/manipulation/planning/world/roboplan_world.py | 2 -- 3 files changed, 3 insertions(+), 12 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index ae3e865c8f..66d33ea359 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -142,10 +142,8 @@ def joint_target_to_global_names( def joint_state_to_ordered_positions( joint_state: JointState, *, - robot_name: str, joint_names: Sequence[str], joint_name_mapping: Mapping[str, str], - context: str, ) -> NDArray[np.float64]: """Convert a JointState to an array ordered by local robot joint names.""" if not joint_state.name: @@ -158,18 +156,15 @@ def joint_state_to_ordered_positions( joint_name_set = set(joint_names) name_to_pos: dict[str, float] = {} - prefix = f"{robot_name}/" for name, position in zip(joint_state.name, joint_state.position, strict=True): if name in joint_name_set: resolved_name = name elif name in joint_name_mapping: resolved_name = joint_name_mapping[name] elif is_global_joint_name(name): - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] + resolved_name = name.split("/", maxsplit=1)[1] if resolved_name not in joint_name_set: - raise ValueError(f"Unknown global joint name for {context}: {name}") + raise ValueError(f"Unknown global joint name: {name}") else: resolved_name = joint_name_mapping.get(name, name) @@ -179,5 +174,5 @@ def joint_state_to_ordered_positions( missing = [name for name in joint_names if name not in name_to_pos] if missing: - raise ValueError(f"JointState missing joints for {context}: {missing}") + raise ValueError(f"JointState missing joints: {missing}") return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index e5cb3cb72a..172e7f40a1 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -876,10 +876,8 @@ def _joint_state_to_q( robot_data = self._robots[robot_id] return joint_state_to_ordered_positions( joint_state, - robot_name=robot_data.config.name, joint_names=robot_data.config.joint_names, joint_name_mapping=robot_data.config.joint_name_mapping, - context="DrakeWorld", ) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 5038fdeebe..b5ac50d828 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -625,10 +625,8 @@ def _joint_state_to_q( robot = self._get_robot(robot_id) return joint_state_to_ordered_positions( joint_state, - robot_name=robot.config.name, joint_names=robot.config.joint_names, joint_name_mapping=robot.config.joint_name_mapping, - context="RoboPlanWorld", ) def _require_finalized(self) -> None: From 3189191b78983b933c25d3a6ff0aa44220a0f56b Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 13:09:13 -0700 Subject: [PATCH 12/29] spec: archive 2 --- .../.openspec.yaml | 2 -- .../planning-world-monitor-groups/README.md | 3 --- .../planning-world-monitor-groups/design.md | 27 ------------------- .../planning-world-monitor-groups/proposal.md | 26 ------------------ .../planning-world-monitor-groups/tasks.md | 19 ------------- .../planning-world-monitor-groups/spec.md | 6 ++++- 6 files changed, 5 insertions(+), 78 deletions(-) delete mode 100644 openspec/changes/planning-world-monitor-groups/.openspec.yaml delete mode 100644 openspec/changes/planning-world-monitor-groups/README.md delete mode 100644 openspec/changes/planning-world-monitor-groups/design.md delete mode 100644 openspec/changes/planning-world-monitor-groups/proposal.md delete mode 100644 openspec/changes/planning-world-monitor-groups/tasks.md rename openspec/{changes/planning-world-monitor-groups => }/specs/planning-world-monitor-groups/spec.md (92%) diff --git a/openspec/changes/planning-world-monitor-groups/.openspec.yaml b/openspec/changes/planning-world-monitor-groups/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/planning-world-monitor-groups/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/planning-world-monitor-groups/README.md b/openspec/changes/planning-world-monitor-groups/README.md deleted file mode 100644 index c5b27282f4..0000000000 --- a/openspec/changes/planning-world-monitor-groups/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# planning-world-monitor-groups - -Extract PR 2 from reference branch cc/spec/movegroup: teach world backends and world monitors to answer planning-group FK, Jacobian, state, and collision queries. diff --git a/openspec/changes/planning-world-monitor-groups/design.md b/openspec/changes/planning-world-monitor-groups/design.md deleted file mode 100644 index eb1e36a5e6..0000000000 --- a/openspec/changes/planning-world-monitor-groups/design.md +++ /dev/null @@ -1,27 +0,0 @@ -## Context - -Planning groups become useful when world backends can answer queries for a specific group. This layer should prove that configured groups map correctly onto backend model instances, joint order, FK target frames, and Jacobian columns. - -## Goals / Non-Goals - -**Goals:** -- Implement group FK/Jacobian support in Drake and RoboPlan worlds. -- Make WorldMonitor expose group-scoped state and query helpers. -- Validate group-local joint ordering and global/local joint-state mapping. -- Keep base pose/base link semantics safe by rejecting unsupported backend cases rather than silently producing wrong transforms. - -**Non-Goals:** -- Do not migrate planner or IK algorithms yet. -- Do not expose new public module APIs yet. -- Do not include Viser UI changes. - -## Decisions - -- Robot-scoped FK/Jacobian compatibility wrappers may remain temporarily but must resolve through exactly one pose-targetable group or fail clearly. -- RoboPlan Jacobians must be projected into group-local joint order; returning a full backend Jacobian is not sufficient. -- Backends should reject unsupported base-pose/base-link combinations explicitly. - -## Risks / Trade-offs - -- Backend capabilities differ. Tests should pin the supported subset and error messages. -- Some existing tests may need compatibility shims until PR 3 migrates algorithms and module APIs. diff --git a/openspec/changes/planning-world-monitor-groups/proposal.md b/openspec/changes/planning-world-monitor-groups/proposal.md deleted file mode 100644 index 2d8d1681d1..0000000000 --- a/openspec/changes/planning-world-monitor-groups/proposal.md +++ /dev/null @@ -1,26 +0,0 @@ -## Why - -Once planning groups exist, the backend query layer must answer group-scoped FK, Jacobian, joint-state, and collision questions. This is the second reviewable PR in the stack: it consumes the foundation without exposing the public `ManipulationModule` API or UI rewrite. - -## What Changes - -- Teach Drake and RoboPlan world backends to resolve planning groups. -- Add group FK/Jacobian APIs and group-local joint ordering behavior. -- Update world monitor and robot state monitor logic for group-aware state access. -- Preserve legacy robot-scoped wrappers only where needed for compatibility, resolving through a unique pose-targetable group. -- Add backend and monitor tests. - -## Capabilities - -### New Capabilities -- `planning-world-monitor-groups`: World and monitor support for planning-group queries. - -### Modified Capabilities -- `manipulation-world-backends`: Drake/RoboPlan backends understand group-local chains and target frames. - -## Impact - -- Base branch: PR 1 `planning-groups-foundation`. -- Reference implementation: `cc/spec/movegroup`. -- Primary files: `dimos/manipulation/planning/world/*`, `dimos/manipulation/planning/monitor/*`, RoboPlan world tests, Drake planning-group tests. -- Out of scope: IK/RRT algorithm migration, `ManipulationModule`, Viser UI, and control integration. diff --git a/openspec/changes/planning-world-monitor-groups/tasks.md b/openspec/changes/planning-world-monitor-groups/tasks.md deleted file mode 100644 index 8e31566a05..0000000000 --- a/openspec/changes/planning-world-monitor-groups/tasks.md +++ /dev/null @@ -1,19 +0,0 @@ -## 1. Extract backend group support - -- [x] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. -- [x] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. -- [x] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. -- [x] 1.4 Extract world monitor and robot state monitor group-aware query changes. - -## 2. Tests - -- [x] 2.1 Bring over Drake group world tests. -- [x] 2.2 Bring over RoboPlan world tests, including planning split file if needed. -- [x] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. - -## 3. Validation - -- [x] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. - - Note: `dimos/manipulation/test_roboplan_world_planning.py` does not exist in this branch; ran the same targeted command without that optional split file. -- [x] 3.2 Run targeted mypy on changed backend/monitor files. -- [x] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. diff --git a/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md b/openspec/specs/planning-world-monitor-groups/spec.md similarity index 92% rename from openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md rename to openspec/specs/planning-world-monitor-groups/spec.md index 3bff6f8bd0..5a737525bc 100644 --- a/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md +++ b/openspec/specs/planning-world-monitor-groups/spec.md @@ -1,4 +1,8 @@ -## ADDED Requirements +## Purpose + +Specify world and monitor support for planning-group queries. + +## Requirements ### Requirement: World backends must answer group pose queries World backends MUST provide end-effector pose queries for pose-targetable planning groups using the group's configured target frame. From 8a4ddbdaf5b13c1e3b4d365fa472d62fe4514116 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 13:43:39 -0700 Subject: [PATCH 13/29] test: improve planning group coverage --- .../planning/groups/test_planning_groups.py | 105 +++++++++++++ .../planning/monitor/test_world_monitor.py | 81 ++++++++++ .../planning/world/drake_world.py | 9 +- .../world/test_drake_world_planning_groups.py | 50 +++++- dimos/manipulation/test_roboplan_world.py | 142 ++++++++++++++++++ 5 files changed, 380 insertions(+), 7 deletions(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index e716366202..0d5f90a2c9 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -32,6 +32,7 @@ from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.groups.utils import ( filter_joint_state_to_selected_joints, + joint_state_to_ordered_positions, joint_target_to_global_names, matching_global_joint_name, planning_group_id_from_selector, @@ -473,6 +474,37 @@ def test_matching_global_joint_name_requires_unique_suffix_match() -> None: assert matching_global_joint_name({"left/j1": 1.0}, "j2") is None +def test_filter_joint_state_to_selected_joints_uses_local_fallbacks() -> None: + state = JointState(name=["j2", "arm/j1"], position=[2.0, 1.0]) + + filtered = filter_joint_state_to_selected_joints( + state, + ["arm/j1", "arm/j2"], + ["j1", "j2"], + ) + + assert filtered.name == ["arm/j1", "arm/j2"] + assert filtered.position == [1.0, 2.0] + + +def test_joint_target_to_global_names_accepts_unnamed_positions_in_group_order() -> None: + target = joint_target_to_global_names( + PlanningGroup( + id="left/arm", + robot_name="left", + group_name="arm", + joint_names=("left/j2", "left/j1"), + local_joint_names=("j2", "j1"), + base_link="base", + tip_link="ee", + ), + JointState(name=[], position=[2.0, 1.0]), + ) + + assert target.name == ["left/j2", "left/j1"] + assert target.position == [2.0, 1.0] + + def test_planning_group_id_from_selector_accepts_id_or_group() -> None: group = _make_group() @@ -551,3 +583,76 @@ def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> with pytest.raises(ValueError, match="multiple pose-target planning groups"): _ = config.end_effector_link + + +def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: + joint_names = ["joint1", "joint2", "joint3"] + mapping = {"hw1": "joint1", "hw2": "joint2", "hw3": "joint3"} + + unnamed = joint_state_to_ordered_positions( + JointState(name=[], position=[1.0, 2.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + local = joint_state_to_ordered_positions( + JointState(name=["joint3", "joint1", "joint2"], position=[30.0, 10.0, 20.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + coordinator = joint_state_to_ordered_positions( + JointState(name=["hw2", "hw3", "hw1"], position=[200.0, 300.0, 100.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + global_names = joint_state_to_ordered_positions( + JointState(name=["arm/joint2", "arm/joint1", "arm/joint3"], position=[2.0, 1.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + + assert unnamed.tolist() == [1.0, 2.0, 3.0] + assert local.tolist() == [10.0, 20.0, 30.0] + assert coordinator.tolist() == [100.0, 200.0, 300.0] + assert global_names.tolist() == [1.0, 2.0, 3.0] + + +def test_joint_state_to_ordered_positions_rejects_invalid_inputs() -> None: + joint_names = ["joint1", "joint2"] + mapping = {"hw1": "joint1"} + + with pytest.raises(ValueError, match="position length"): + joint_state_to_ordered_positions( + JointState(name=[], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="name and position"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "joint2"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="duplicate"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "hw1"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="Unknown global"): + joint_state_to_ordered_positions( + JointState(name=["arm/joint3", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="missing joints"): + joint_state_to_ordered_positions( + JointState(name=["joint1"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="missing joints"): + joint_state_to_ordered_positions( + JointState(name=["mystery", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 84a17160d1..a78957858c 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -340,6 +340,87 @@ def test_current_group_joint_state_uses_public_names_in_group_order() -> None: assert state.position == [0.2, 0.1] +def test_current_global_joint_state_skips_stale_robots_and_preserves_state_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + fresh_id = monitor.add_robot(_three_joint_reordered_group_config()) + stale_id = monitor.add_robot( + RobotModelConfig( + name="arm2", + model_path=Path("/tmp/arm2.urdf"), + joint_names=["a", "b"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("a", "b"), base_link="base", tip_link="ee" + ) + ], + ) + ) + monitor._state_monitors[fresh_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + monitor._state_monitors[stale_id] = _FakeStateMonitor([1.0, 2.0], stale=True) # type: ignore[attr-defined] + monitor.add_robot( + RobotModelConfig( + name="arm3", + model_path=Path("/tmp/arm3.urdf"), + joint_names=["x"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("x",), base_link="base", tip_link="ee" + ) + ], + ) + ) + + state = monitor.current_global_joint_state(max_age=0.5) + + assert state.name == ["arm/j1", "arm/j2", "arm/j3"] + assert state.position == [0.1, 0.2, 0.3] + + +def test_current_group_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.current_group_joint_state("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.current_group_joint_state("arm/manipulator") + + +def test_group_ee_pose_uses_current_state_when_no_joint_state_is_provided() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + pose = monitor.get_group_ee_pose("arm/manipulator") + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + +def test_group_ee_pose_without_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.get_group_ee_pose("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.get_group_ee_pose("arm/manipulator") + + def test_group_kinematics_with_full_state_does_not_require_current_state() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 172e7f40a1..dae0d3f73c 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -334,7 +334,8 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config - def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + @staticmethod + def _validate_planning_group_config(config: RobotModelConfig) -> None: seen_group_names: set[str] = set() for definition in config.planning_groups: make_planning_group_id(config.name, definition.name) @@ -343,8 +344,9 @@ def _validate_planning_group_config(self, config: RobotModelConfig) -> None: make_global_joint_names(config.name, definition.joint_names) seen_group_names.add(definition.name) + @staticmethod def _planning_group_from_config( - self, config: RobotModelConfig, group_id: PlanningGroupID + config: RobotModelConfig, group_id: PlanningGroupID ) -> PlanningGroup: for definition in config.planning_groups: if make_planning_group_id(config.name, definition.name) == group_id: @@ -369,7 +371,8 @@ def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: continue raise KeyError(f"Unknown planning group ID: {group_id}") - def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + @staticmethod + def _primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: pose_groups = [group for group in config.planning_groups if group.has_pose_target] if not pose_groups: return None diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index 1331e447b7..c981c76896 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -19,15 +19,17 @@ import numpy as np import pytest -pydrake = pytest.importorskip("pydrake") -_ = pydrake - from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.world.drake_world import DrakeWorld +from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +requires_drake = pytest.mark.skipif( + not DRAKE_AVAILABLE, + reason="Drake planning-group tests require the manipulation extra", +) + def _write_urdf(path: Path) -> None: path.write_text( @@ -72,6 +74,43 @@ def _arm_group( ) +def test_drake_config_group_helpers_resolve_groups_without_drake_runtime(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + config = _config(urdf, [_arm_group("joint2", "joint1", name="wrist")]) + + group = DrakeWorld._planning_group_from_config(config, "arm/wrist") + + assert DrakeWorld._primary_pose_group_id_for_config(config) == "arm/wrist" + assert group.id == "arm/wrist" + assert group.joint_names == ("arm/joint2", "arm/joint1") + assert group.local_joint_names == ("joint2", "joint1") + assert group.tip_link == "tool0" + + +def test_drake_config_group_helpers_validate_duplicate_and_ambiguous_groups( + tmp_path: Path, +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + duplicate = _config( + urdf, + [_arm_group("joint1", name="same"), _arm_group("joint2", name="same")], + ) + ambiguous = _config( + urdf, + [_arm_group("joint1", name="a"), _arm_group("joint2", name="b")], + ) + + with pytest.raises(ValueError, match="already registered"): + DrakeWorld._validate_planning_group_config(duplicate) + with pytest.raises(ValueError, match="multiple pose"): + DrakeWorld._primary_pose_group_id_for_config(ambiguous) + with pytest.raises(KeyError, match="Unknown planning group ID"): + DrakeWorld._planning_group_from_config(ambiguous, "arm/missing") + + +@requires_drake def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -91,6 +130,7 @@ def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Pat assert world.get_jacobian(ctx, robot_id).shape == (6, 2) +@requires_drake def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -118,6 +158,7 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non np.testing.assert_allclose(reverse_jacobian[:, 1], forward_jacobian[:, 0]) +@requires_drake def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -142,6 +183,7 @@ def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_pa ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) +@requires_drake def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index f3879297d5..32eaa97bfa 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -563,6 +563,148 @@ def fake_jacobian( np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 1]]) +def test_group_kinematics_reject_missing_tip_or_missing_context( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_tip_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="joint_only", joint_names=("joint1", "joint2"), base_link="base" + ) + ] + } + ) + world, robot_id = _make_world(fake_roboplan, no_tip_config) + world.finalize() + + with pytest.raises(ValueError, match="no tip link"): + world.get_group_ee_pose(world.get_live_context(), "arm/joint_only") + with pytest.raises(ValueError, match="no tip link"): + world.get_group_jacobian(world.get_live_context(), "arm/joint_only") + + ctx = world.get_live_context() + del ctx.q_by_robot[robot_id] + with pytest.raises(KeyError, match=robot_id): + world.get_link_pose(ctx, robot_id, "tcp") + + jacobian_world, jacobian_robot_id = _make_world(fake_roboplan, robot_config) + jacobian_world.finalize() + jacobian_ctx = jacobian_world.get_live_context() + del jacobian_ctx.q_by_robot[jacobian_robot_id] + with pytest.raises(KeyError, match=jacobian_robot_id): + jacobian_world.get_group_jacobian(jacobian_ctx, "arm/manipulator") + + +def test_group_jacobian_validates_projection_shape_and_joint_names( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "other"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 2)), + ) + with pytest.raises(ValueError, match="Unknown joints"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "joint2"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((5, 2)), + ) + with pytest.raises(ValueError, match="Unexpected RoboPlan Jacobian shape"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 3)), + ) + with pytest.raises(ValueError, match="cannot project"): + world.get_group_jacobian(ctx, "arm/manipulator") + + +def test_group_jacobian_falls_back_to_configured_joint_order_when_scene_order_is_unavailable( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + def missing_group_info(self: FakeScene, name: str) -> FakeJointGroupInfo: + _ = (self, name) + raise AttributeError("no joint group info") + + monkeypatch.setattr(FakeScene, "getJointGroupInfo", missing_group_info) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + np.testing.assert_allclose( + world.get_group_jacobian(ctx, "arm/manipulator"), + np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + +def test_legacy_kinematics_wrappers_require_unique_pose_group( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_pose_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition(name="base", joint_names=("joint1",), base_link="base") + ] + } + ) + no_pose_world, no_pose_id = _make_world(fake_roboplan, no_pose_config) + no_pose_world.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_ee_pose(no_pose_world.get_live_context(), no_pose_id) + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_jacobian(no_pose_world.get_live_context(), no_pose_id) + + ambiguous_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="a", joint_names=("joint1",), base_link="base", tip_link="a_tip" + ), + PlanningGroupDefinition( + name="b", joint_names=("joint2",), base_link="base", tip_link="b_tip" + ), + ] + } + ) + ambiguous_world, ambiguous_id = _make_world(fake_roboplan, ambiguous_config) + ambiguous_world.finalize() + with pytest.raises(ValueError, match="pose-targetable planning groups"): + ambiguous_world.get_jacobian(ambiguous_world.get_live_context(), ambiguous_id) + + +def test_group_lookup_rejects_unknown_group_id( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(KeyError, match="Unknown planning group ID"): + world.get_group_ee_pose(world.get_live_context(), "other/missing") + + def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world, robot_id = _make_world(fake_roboplan, robot_config) world.finalize() From 5c3893567bf0d1799231b2f9244bf901208053f6 Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:45:50 -0700 Subject: [PATCH 14/29] Update dimos/manipulation/planning/groups/utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- dimos/manipulation/planning/groups/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 66d33ea359..31130084de 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -166,7 +166,7 @@ def joint_state_to_ordered_positions( if resolved_name not in joint_name_set: raise ValueError(f"Unknown global joint name: {name}") else: - resolved_name = joint_name_mapping.get(name, name) + raise ValueError(f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name") if resolved_name in name_to_pos: raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") From 562f86ae98acfea7242d2ddaa79eb94e25e80b65 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 20:46:27 +0000 Subject: [PATCH 15/29] [autofix.ci] apply automated fixes --- dimos/manipulation/planning/groups/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 31130084de..283681adcc 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -166,7 +166,9 @@ def joint_state_to_ordered_positions( if resolved_name not in joint_name_set: raise ValueError(f"Unknown global joint name: {name}") else: - raise ValueError(f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name") + raise ValueError( + f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name" + ) if resolved_name in name_to_pos: raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") From 90baf52f582de827555436ce82b8e5f138e412a3 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 14:10:26 -0700 Subject: [PATCH 16/29] test: fix joint state ordering assertion --- dimos/manipulation/planning/groups/test_planning_groups.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 0d5f90a2c9..8b399354a6 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -650,7 +650,7 @@ def test_joint_state_to_ordered_positions_rejects_invalid_inputs() -> None: joint_names=joint_names, joint_name_mapping=mapping, ) - with pytest.raises(ValueError, match="missing joints"): + with pytest.raises(ValueError, match="Unrecognized joint name"): joint_state_to_ordered_positions( JointState(name=["mystery", "joint2"], position=[1.0, 2.0]), joint_names=joint_names, From f3aac24b2e3d26f1a1a23d36096bc88a4cd4cac5 Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:57:49 -0700 Subject: [PATCH 17/29] Planning Group 3: Add group-aware IK and RRT planning (#2747) --- .../kinematics/drake_optimization_ik.py | 129 +++++- .../planning/kinematics/jacobian_ik.py | 115 +++++- .../planning/kinematics/pink_ik.py | 347 ++++++++++++++++- .../test_drake_optimization_ik_selection.py | 253 ++++++++++++ .../kinematics/test_jacobian_ik_selection.py | 147 +++++++ .../planning/kinematics/test_pink_ik.py | 366 +++++++++++++++++- .../manipulation/planning/kinematics/utils.py | 236 +++++++++++ .../planning/planners/rrt_planner.py | 197 +++++++++- .../planning/planners/selected_joint_space.py | 319 +++++++++++++++ .../planners/test_rrt_planner_selection.py | 199 ++++++++++ dimos/manipulation/planning/spec/enums.py | 2 + dimos/manipulation/planning/spec/protocols.py | 28 ++ .../planning/world/roboplan_world.py | 17 +- .../changes/group-aware-ik-rrt/.openspec.yaml | 2 - openspec/changes/group-aware-ik-rrt/README.md | 3 - openspec/changes/group-aware-ik-rrt/design.md | 27 -- .../changes/group-aware-ik-rrt/proposal.md | 24 -- .../specs/group-aware-ik-rrt/spec.md | 22 -- openspec/changes/group-aware-ik-rrt/tasks.md | 19 - openspec/specs/group-aware-ik-rrt/spec.md | 38 ++ 20 files changed, 2338 insertions(+), 152 deletions(-) create mode 100644 dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py create mode 100644 dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py create mode 100644 dimos/manipulation/planning/kinematics/utils.py create mode 100644 dimos/manipulation/planning/planners/selected_joint_space.py create mode 100644 dimos/manipulation/planning/planners/test_rrt_planner_selection.py delete mode 100644 openspec/changes/group-aware-ik-rrt/.openspec.yaml delete mode 100644 openspec/changes/group-aware-ik-rrt/README.md delete mode 100644 openspec/changes/group-aware-ik-rrt/design.md delete mode 100644 openspec/changes/group-aware-ik-rrt/proposal.md delete mode 100644 openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md delete mode 100644 openspec/changes/group-aware-ik-rrt/tasks.md create mode 100644 openspec/specs/group-aware-ik-rrt/spec.md diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index abccee119d..91da25986c 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -16,10 +16,17 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics.utils import ( + filter_result_to_group as _filter_result_to_group, + resolve_single_pose_target_request as _resolve_single_pose_target_request, + unique_pose_target_frame_for_robot as _unique_pose_target_frame_for_robot, +) from dimos.manipulation.planning.spec.enums import IKStatus from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec @@ -28,7 +35,6 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray @@ -84,6 +90,13 @@ def solve( if error is not None: return error + target_frame_name = _unique_pose_target_frame_for_robot(world, robot_id) + if target_frame_name is None: + return _create_failure_result( + IKStatus.UNSUPPORTED, + "DrakeOptimizationIK requires exactly one pose-targetable planning group for legacy solve()", + ) + # Convert PoseStamped to 4x4 matrix via Transform target_matrix = Transform( translation=target_pose.position, @@ -127,6 +140,7 @@ def solve( orientation_tolerance=orientation_tolerance, lower_limits=lower_limits, upper_limits=upper_limits, + target_frame_name=target_frame_name, ) if result.is_success() and result.joint_state is not None: @@ -156,6 +170,96 @@ def solve( f"IK failed after {max_attempts} attempts", ) + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve a planning-group-scoped pose target with Drake IK.""" + error = self._validate_world(world) + if error is not None: + return error + request, request_error = _resolve_single_pose_target_request( + world, + pose_targets, + auxiliary_groups, + seed, + "DrakeOptimizationIK", + ) + if request_error is not None: + return request_error + if request is None or request.group.tip_link is None: + return _create_failure_result( + IKStatus.UNSUPPORTED, + "DrakeOptimizationIK requires a pose-targetable planning group", + ) + + lower_limits, upper_limits = world.get_joint_limits(request.robot_id) + target_matrix = Transform( + translation=request.target_pose.position, + rotation=request.target_pose.orientation, + ).to_matrix() + target_transform = RigidTransform(target_matrix) + locked_positions = { + index: float(request.seed_positions[index]) + for index in range(len(request.joint_names)) + if index not in set(request.group_indices) + } + + best_result: IKResult | None = None + best_error = float("inf") + for attempt in range(max_attempts): + if attempt == 0: + current_seed = request.seed_positions + else: + current_seed = request.seed_positions.copy() + random_group_positions = np.random.uniform( + lower_limits[request.group_indices], upper_limits[request.group_indices] + ) + current_seed[request.group_indices] = random_group_positions + + result = self._solve_single( + world=world, + robot_id=request.robot_id, + target_transform=target_transform, + seed=current_seed, + joint_names=request.joint_names, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + lower_limits=lower_limits, + upper_limits=upper_limits, + target_frame_name=request.group.tip_link, + locked_joint_positions=locked_positions, + ) + if not result.is_success() or result.joint_state is None: + continue + if check_collision and not world.check_config_collision_free( + request.robot_id, result.joint_state + ): + continue + total_error = result.position_error + result.orientation_error + if total_error < best_error: + best_error = total_error + best_result = result + if ( + result.position_error <= position_tolerance + and result.orientation_error <= orientation_tolerance + ): + return _filter_result_to_group(result, request.group) + + if best_result is not None: + return _filter_result_to_group(best_result, request.group) + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"IK failed after {max_attempts} attempts", + ) + def _solve_single( self, world: WorldSpec, @@ -167,6 +271,8 @@ def _solve_single( orientation_tolerance: float, lower_limits: NDArray[np.float64], upper_limits: NDArray[np.float64], + target_frame_name: str, + locked_joint_positions: Mapping[int, float] | None = None, ) -> IKResult: # Get robot data from world internals (Drake-specific access) robot_data = world._robots[robot_id] # type: ignore[attr-defined] @@ -175,12 +281,13 @@ def _solve_single( # Create IK problem ik = InverseKinematics(plant) - # Get end-effector frame - ee_frame = robot_data.ee_frame + target_frame = plant.GetBodyByName( + target_frame_name, robot_data.model_instance + ).body_frame() # Add position constraint ik.AddPositionConstraint( - frameB=ee_frame, + frameB=target_frame, p_BQ=np.array([0.0, 0.0, 0.0]), # type: ignore[arg-type] frameA=plant.world_frame(), p_AQ_lower=target_transform.translation() - np.array([position_tolerance] * 3), @@ -191,7 +298,7 @@ def _solve_single( ik.AddOrientationConstraint( frameAbar=plant.world_frame(), R_AbarA=target_transform.rotation(), - frameBbar=ee_frame, + frameBbar=target_frame, R_BbarB=RotationMatrix(), theta_bound=orientation_tolerance, ) @@ -200,6 +307,10 @@ def _solve_single( prog = ik.get_mutable_prog() q = ik.q() + for local_index, value in (locked_joint_positions or {}).items(): + joint_idx = robot_data.joint_indices[local_index] + prog.AddBoundingBoxConstraint(value, value, q[joint_idx]) + # Set initial guess (full positions vector) full_seed = np.zeros(plant.num_positions()) for i, joint_idx in enumerate(robot_data.joint_indices): @@ -223,13 +334,13 @@ def _solve_single( joint_solution = np.clip(joint_solution, lower_limits, upper_limits) # Compute actual error using FK - solution_state = JointState(name=joint_names, position=joint_solution.tolist()) + solution_state = JointState({"name": joint_names, "position": joint_solution.tolist()}) with world.scratch_context() as ctx: world.set_joint_state(ctx, robot_id, solution_state) - actual_pose = world.get_ee_pose(ctx, robot_id) + actual_matrix = world.get_link_pose(ctx, robot_id, target_frame_name) position_error, orientation_error = compute_pose_error( - pose_to_matrix(actual_pose), + actual_matrix, target_transform.GetAsMatrix4(), # type: ignore[arg-type] ) @@ -251,7 +362,7 @@ def _create_success_result( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index 7727b6fa0f..4c4e16207a 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -24,10 +24,15 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics.utils import ( + resolve_single_pose_target_request as _resolve_single_pose_target_request, +) from dimos.manipulation.planning.spec.enums import IKStatus from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec @@ -145,7 +150,9 @@ def solve( else: # Random seed within joint limits random_positions = np.random.uniform(lower_limits, upper_limits) - current_seed = JointState(name=joint_names, position=random_positions.tolist()) + current_seed = JointState( + {"name": joint_names, "position": random_positions.tolist()} + ) # Solve iterative IK result = self.solve_iterative( @@ -185,6 +192,61 @@ def solve( f"IK failed after {max_attempts} attempts", ) + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve a planning-group pose target using group FK/Jacobian.""" + if not world.is_finalized: + return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") + request, request_error = _resolve_single_pose_target_request( + world, + pose_targets, + auxiliary_groups, + seed, + "JacobianIK", + ) + if request_error is not None: + return request_error + if request is None: + return _create_failure_result(IKStatus.NO_SOLUTION, "Invalid pose target request") + + full_seed = JointState( + {"name": request.joint_names, "position": request.seed_positions.tolist()} + ) + result = self.solve_iterative( + world=world, + robot_id=request.robot_id, + target_pose=request.target_pose, + seed=full_seed, + max_iterations=self._max_iterations * max(1, max_attempts), + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + group=request.group, + active_joint_indices=request.group_indices, + ) + if not result.is_success() or result.joint_state is None: + return result + + if check_collision: + full_positions = request.seed_positions.copy() + full_positions[request.group_indices] = np.asarray( + result.joint_state.position, dtype=np.float64 + ) + full_state = JointState( + {"name": request.joint_names, "position": full_positions.tolist()} + ) + if not world.check_config_collision_free(request.robot_id, full_state): + return _create_failure_result(IKStatus.COLLISION, "IK solution is in collision") + return result + def solve_iterative( self, world: WorldSpec, @@ -194,6 +256,8 @@ def solve_iterative( max_iterations: int = 100, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, + group: PlanningGroup | None = None, + active_joint_indices: list[int] | None = None, ) -> IKResult: """Iterative Jacobian-based IK until convergence. @@ -219,6 +283,8 @@ def solve_iterative( ).to_matrix() current_joints = np.array(seed.position, dtype=np.float64) joint_names = seed.name + active_joint_indices = active_joint_indices or list(range(len(joint_names))) + result_joint_names = list(group.joint_names) if group is not None else joint_names max_iterations = max_iterations or self._max_iterations lower_limits, upper_limits = world.get_joint_limits(robot_id) @@ -226,11 +292,15 @@ def solve_iterative( for iteration in range(max_iterations): with world.scratch_context() as ctx: # Set current position (convert to JointState for API) - current_state = JointState(name=joint_names, position=current_joints.tolist()) + current_state = JointState( + {"name": joint_names, "position": current_joints.tolist()} + ) world.set_joint_state(ctx, robot_id, current_state) - # Get current pose (as matrix for error computation) - current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + if group is None: + current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + else: + current_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) # Compute error pos_error, ori_error = compute_pose_error(current_pose, target_matrix) @@ -238,28 +308,30 @@ def solve_iterative( # Check convergence if pos_error <= position_tolerance and ori_error <= orientation_tolerance: return _create_success_result( - joint_names=joint_names, - joint_positions=current_joints, + joint_names=result_joint_names, + joint_positions=current_joints[active_joint_indices], position_error=pos_error, orientation_error=ori_error, iterations=iteration + 1, ) + if group is None: + jacobian = world.get_jacobian(ctx, robot_id) + else: + jacobian = world.get_group_jacobian(ctx, group.id) + # Compute twist to reduce error twist = compute_error_twist(current_pose, target_matrix, gain=0.5) - # Get Jacobian - J = world.get_jacobian(ctx, robot_id) - # Adaptive damping near singularities - if check_singularity(J, threshold=self._singularity_threshold): + if check_singularity(jacobian, threshold=self._singularity_threshold): # Increase damping near singularity instead of failing effective_damping = self._damping * 10.0 else: effective_damping = self._damping # Compute joint velocities - J_pinv = damped_pseudoinverse(J, effective_damping) + J_pinv = damped_pseudoinverse(jacobian, effective_damping) q_dot = J_pinv @ twist # Clamp maximum joint change per iteration (like reference implementations) @@ -268,16 +340,23 @@ def solve_iterative( if max_change > max_delta: q_dot = q_dot * (max_delta / max_change) - current_joints = current_joints + q_dot + current_joints[active_joint_indices] = current_joints[active_joint_indices] + q_dot # Clip to limits - current_joints = np.clip(current_joints, lower_limits, upper_limits) + current_joints[active_joint_indices] = np.clip( + current_joints[active_joint_indices], + lower_limits[active_joint_indices], + upper_limits[active_joint_indices], + ) # Compute final error with world.scratch_context() as ctx: - final_state = JointState(name=joint_names, position=current_joints.tolist()) + final_state = JointState({"name": joint_names, "position": current_joints.tolist()}) world.set_joint_state(ctx, robot_id, final_state) - final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + if group is None: + final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + else: + final_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) pos_error, ori_error = compute_pose_error(final_pose, target_matrix) return _create_failure_result( @@ -349,7 +428,7 @@ def solve_differential( if max_ratio > 1.0: q_dot = q_dot / max_ratio - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState({"name": joint_names, "velocity": q_dot.tolist()}) def solve_differential_position_only( self, @@ -397,7 +476,7 @@ def solve_differential_position_only( # Compute joint velocities q_dot = J_pinv @ vel_array - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState({"name": joint_names, "velocity": q_dot.tolist()}) # Result Helpers @@ -413,7 +492,7 @@ def _create_success_result( """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index 1245e5aea4..3456c2e27f 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from dataclasses import dataclass import importlib from pathlib import Path @@ -24,10 +25,17 @@ import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig +from dimos.manipulation.planning.kinematics.utils import ( + groups_by_robot as _groups_by_robot, + robot_ids_by_name as _robot_ids_by_name, + seed_positions_with_world_fallback as _seed_positions_with_world_fallback, + unique_pose_target_frame_for_robot as _unique_pose_target_frame_for_robot, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus -from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID +from dimos.manipulation.planning.spec.models import IKResult, RobotName, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake @@ -98,7 +106,7 @@ def __init__( config_values.update(overrides) self.config = PinkKinematicsConfig(**config_values) self._modules = _load_optional_dependencies(self.config.solver) - self._robot_contexts: dict[str, _PinkRobotContext] = {} + self._robot_contexts: dict[tuple[str, str], _PinkRobotContext] = {} def solve( self, @@ -115,8 +123,15 @@ def solve( if not world.is_finalized: return _failure(IKStatus.NO_SOLUTION, "World must be finalized before IK") + target_frame_name = _unique_pose_target_frame_for_robot(world, robot_id) + if target_frame_name is None: + return _failure( + IKStatus.NO_SOLUTION, + "PinkIK requires exactly one pose-targetable planning group for legacy solve()", + ) + try: - robot_context = self._get_robot_context(world, robot_id) + robot_context = self._get_robot_context(world, robot_id, target_frame_name) except (FileNotFoundError, ImportError, ValueError) as exc: return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") @@ -164,6 +179,265 @@ def solve( return _failure(IKStatus.NO_SOLUTION, f"Pink IK failed after {max_attempts} attempts") + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve planning-group-scoped pose targets with Pink IK.""" + if not world.is_finalized: + return _failure(IKStatus.NO_SOLUTION, "World must be finalized before IK") + all_groups = tuple(pose_targets.keys()) + tuple(auxiliary_groups) + if not all_groups: + return _failure( + IKStatus.NO_SOLUTION, "At least one pose target or auxiliary group is required" + ) + bad_groups = [ + group.id + for group in pose_targets + if not group.has_pose_target or group.tip_link is None + ] + if bad_groups: + return _failure( + IKStatus.UNSUPPORTED, + f"Planning groups have no pose target frame: {bad_groups}", + ) + + try: + selection = PlanningGroupSelection.from_groups(all_groups) + robot_ids_by_name = _robot_ids_by_name(world, selection.robot_names) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, str(exc)) + + results_by_robot: dict[RobotName, IKResult] = {} + for robot_name, groups in _groups_by_robot(all_groups).items(): + robot_id = robot_ids_by_name[robot_name] + config = world.get_robot_config(robot_id) + joint_names = list(config.joint_names) + try: + selected_indices = [ + joint_names.index(name) for group in groups for name in group.local_joint_names + ] + seed_positions = _seed_positions_with_world_fallback( + world, robot_id, config.name, joint_names, seed + ) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK mapping failed: {exc}") + robot_pose_targets = [group for group in groups if group in pose_targets] + if not robot_pose_targets: + robot_result = _success(joint_names, seed_positions, 0.0, 0.0, 0) + results_by_robot[robot_name] = robot_result + continue + + lower_limits, upper_limits = world.get_joint_limits(robot_id) + locked_positions = { + index: float(seed_positions[index]) + for index in range(len(joint_names)) + if index not in set(selected_indices) + } + targets: list[tuple[_PinkRobotContext, NDArray[np.float64]]] = [] + try: + for group in robot_pose_targets: + if group.tip_link is None: + raise ValueError(f"Planning group '{group.id}' has no pose target frame") + targets.append( + ( + self._get_robot_context(world, robot_id, group.tip_link), + self._target_in_model_frame(config, pose_targets[group]), + ) + ) + except (FileNotFoundError, ImportError, ValueError) as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") + + fallback_result: IKResult | None = None + for attempt in range(max_attempts): + current_positions = seed_positions.copy() + if attempt > 0: + current_positions[selected_indices] = np.random.uniform( + lower_limits[selected_indices], upper_limits[selected_indices] + ) + try: + q0 = self._q_from_dimos_positions(targets[0][0], current_positions) + if len(targets) == 1: + result = self._solve_single( + robot_context=targets[0][0], + target_model=targets[0][1], + seed_q=q0, + lower_limits=lower_limits, + upper_limits=upper_limits, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + locked_joint_positions=locked_positions, + ) + else: + result = self._solve_multi( + targets=targets, + seed_q=q0, + lower_limits=lower_limits, + upper_limits=upper_limits, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + locked_joint_positions=locked_positions, + ) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK mapping failed: {exc}") + except Exception as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK solver failed: {exc}") + + if not result.is_success() or result.joint_state is None: + if fallback_result is None: + fallback_result = result + continue + results_by_robot[robot_name] = result + break + else: + if fallback_result is not None: + return fallback_result + return _failure( + IKStatus.NO_SOLUTION, f"Pink IK failed after {max_attempts} attempts" + ) + + positions_by_robot: dict[RobotName, dict[str, float]] = {} + max_position_error = 0.0 + max_orientation_error = 0.0 + iterations = 0 + for robot_name, result in results_by_robot.items(): + if not result.is_success() or result.joint_state is None: + return result + positions_by_robot[robot_name] = dict( + zip(result.joint_state.name, result.joint_state.position, strict=True) + ) + max_position_error = max(max_position_error, result.position_error) + max_orientation_error = max(max_orientation_error, result.orientation_error) + iterations = max(iterations, result.iterations) + + selected_names: list[str] = [] + selected_positions: list[float] = [] + for group in selection.groups: + robot_positions = positions_by_robot[group.robot_name] + for global_name, local_name in zip( + group.joint_names, + group.local_joint_names, + strict=True, + ): + if global_name in robot_positions: + position = robot_positions[global_name] + elif local_name in robot_positions: + position = robot_positions[local_name] + else: + return _failure( + IKStatus.NO_SOLUTION, + f"Pink IK result is missing selected joint '{global_name}'", + ) + selected_names.append(global_name) + selected_positions.append(float(position)) + + combined = IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + { + "name": selected_names, + "position": selected_positions, + } + ), + position_error=max_position_error, + orientation_error=max_orientation_error, + iterations=iterations, + message="Pink IK solution found", + ) + if check_collision and not _combined_robot_results_collision_free( + world, + robot_ids_by_name, + results_by_robot, + ): + return _collision_failure(combined) + return combined + + def _solve_multi( + self, + targets: Sequence[tuple[_PinkRobotContext, NDArray[np.float64]]], + seed_q: NDArray[np.float64], + lower_limits: NDArray[np.float64], + upper_limits: NDArray[np.float64], + position_tolerance: float, + orientation_tolerance: float, + locked_joint_positions: Mapping[int, float] | None = None, + ) -> IKResult: + robot_context = targets[0][0] + pink = self._modules.pink + pinocchio = self._modules.pinocchio + configuration = pink.Configuration(robot_context.model, robot_context.data, seed_q.copy()) + tasks: list[Any] = [] + for target_context, target_model in targets: + frame_task = pink.tasks.FrameTask( + target_context.frame_name, + position_cost=self.config.position_cost, + orientation_cost=self.config.orientation_cost, + lm_damping=self.config.lm_damping, + gain=self.config.gain, + ) + frame_task.set_target(_matrix_to_se3(pinocchio, target_model)) + tasks.append(frame_task) + if self.config.posture_cost > 0.0: + posture_task = pink.tasks.PostureTask(cost=self.config.posture_cost) + posture_task.set_target_from_configuration(configuration) + tasks.append(posture_task) + final_position_error = float("inf") + final_orientation_error = float("inf") + for iteration in range(self.config.max_iterations): + errors = [ + compute_pose_error(self._current_frame_matrix(ctx, configuration.q), target_model) + for ctx, target_model in targets + ] + final_position_error = max(error[0] for error in errors) + final_orientation_error = max(error[1] for error in errors) + if ( + final_position_error <= position_tolerance + and final_orientation_error <= orientation_tolerance + ): + return _success( + robot_context.mapping.dimos_joint_names, + self._q_to_dimos_positions(robot_context, configuration.q), + final_position_error, + final_orientation_error, + iteration + 1, + ) + velocity = pink.solve_ik( + configuration, + tasks, + self.config.dt, + solver=self.config.solver, + damping=self.config.damping, + safety_break=self.config.safety_break, + ) + configuration.integrate_inplace(velocity, self.config.dt) + for local_index, value in (locked_joint_positions or {}).items(): + configuration.q[robot_context.mapping.idx_q[local_index]] = value + joint_positions = self._q_to_dimos_positions(robot_context, configuration.q) + if not _within_limits(joint_positions, lower_limits, upper_limits): + return IKResult( + status=IKStatus.JOINT_LIMITS, + joint_state=None, + position_error=final_position_error, + orientation_error=final_orientation_error, + iterations=iteration + 1, + message="Pink IK candidate violates DimOS joint limits", + ) + return IKResult( + status=IKStatus.NO_SOLUTION, + joint_state=None, + position_error=final_position_error, + orientation_error=final_orientation_error, + iterations=self.config.max_iterations, + message="Pink IK did not converge within the iteration budget", + ) + def _solve_single( self, robot_context: _PinkRobotContext, @@ -173,6 +447,7 @@ def _solve_single( upper_limits: NDArray[np.float64], position_tolerance: float, orientation_tolerance: float, + locked_joint_positions: Mapping[int, float] | None = None, ) -> IKResult: pink = self._modules.pink pinocchio = self._modules.pinocchio @@ -224,6 +499,8 @@ def _solve_single( safety_break=self.config.safety_break, ) configuration.integrate_inplace(velocity, self.config.dt) + for local_index, value in (locked_joint_positions or {}).items(): + configuration.q[robot_context.mapping.idx_q[local_index]] = value joint_positions = self._q_to_dimos_positions(robot_context, configuration.q) if not _within_limits(joint_positions, lower_limits, upper_limits): @@ -245,15 +522,20 @@ def _solve_single( message="Pink IK did not converge within the iteration budget", ) - def _get_robot_context(self, world: WorldSpec, robot_id: WorldRobotID) -> _PinkRobotContext: - cache_key = str(robot_id) + def _get_robot_context( + self, + world: WorldSpec, + robot_id: WorldRobotID, + frame_name: str, + ) -> _PinkRobotContext: + cache_key = (str(robot_id), frame_name) if cache_key not in self._robot_contexts: self._robot_contexts[cache_key] = self._build_robot_context( - world.get_robot_config(robot_id) + world.get_robot_config(robot_id), frame_name ) return self._robot_contexts[cache_key] - def _build_robot_context(self, config: RobotModelConfig) -> _PinkRobotContext: + def _build_robot_context(self, config: RobotModelConfig, frame_name: str) -> _PinkRobotContext: pinocchio = self._modules.pinocchio model_path = Path(config.model_path).resolve() if not model_path.exists(): @@ -271,13 +553,14 @@ def _build_robot_context(self, config: RobotModelConfig) -> _PinkRobotContext: model = pinocchio.buildModelFromUrdf(str(prepared_path)) data = model.createData() - frame_id = _get_frame_id(model, config.end_effector_link) + _assert_base_link_is_model_root(model, config.base_link) + frame_id = _get_frame_id(model, frame_name) mapping = _build_joint_mapping(model, config) return _PinkRobotContext( model=model, data=data, frame_id=frame_id, - frame_name=config.end_effector_link, + frame_name=frame_name, mapping=mapping, ) @@ -302,6 +585,21 @@ def _initial_q( q[idx_q] = value return q + def _q_from_dimos_positions( + self, + context: _PinkRobotContext, + positions: NDArray[np.float64], + ) -> NDArray[np.float64]: + pinocchio = self._modules.pinocchio + q = np.array(pinocchio.neutral(context.model), dtype=np.float64) + if len(positions) != len(context.mapping.idx_q): + raise ValueError( + f"Seed has {len(positions)} positions for {len(context.mapping.idx_q)} joints" + ) + for value, idx_q in zip(positions, context.mapping.idx_q, strict=True): + q[idx_q] = value + return q + def _q_to_dimos_positions( self, context: _PinkRobotContext, q: NDArray[np.float64] ) -> NDArray[np.float64]: @@ -407,6 +705,18 @@ def _get_frame_id(model: Any, frame_name: str) -> int: return frame_id +def _assert_base_link_is_model_root(model: Any, base_link: str) -> None: + """Validate that the configured base link is fixed at the Pinocchio model root.""" + frame_id = _get_frame_id(model, base_link) + frame = model.frames[frame_id] + parent_joint = int(getattr(frame, "parentJoint", 0)) + if parent_joint != 0: + raise ValueError( + f"PinkIK expects RobotModelConfig.base_link '{base_link}' to be the model root; " + f"Pinocchio frame parentJoint is {parent_joint}" + ) + + def _missing_joint_message(model: Any, joint_name: str) -> str: available = [str(name) for name in getattr(model, "names", [])] return f"Joint '{joint_name}' not found in Pinocchio model. Available joints: {available}" @@ -456,6 +766,23 @@ def _within_limits( ) +def _combined_robot_results_collision_free( + world: WorldSpec, + robot_ids_by_name: Mapping[RobotName, WorldRobotID], + results_by_robot: Mapping[RobotName, IKResult], +) -> bool: + with world.scratch_context() as ctx: + for robot_name, result in results_by_robot.items(): + if result.joint_state is None: + return False + world.set_joint_state(ctx, robot_ids_by_name[robot_name], result.joint_state) + return all( + world.is_collision_free(ctx, robot_id) + for robot_name, robot_id in robot_ids_by_name.items() + if robot_name in results_by_robot + ) + + def _success( joint_names: list[str], joint_positions: NDArray[np.float64], @@ -465,7 +792,7 @@ def _success( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py new file mode 100644 index 0000000000..afac663e72 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py @@ -0,0 +1,253 @@ +# 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. + +from __future__ import annotations + +from contextlib import contextmanager +from pathlib import Path + +import numpy as np + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics import drake_optimization_ik as drake_ik +from dimos.manipulation.planning.kinematics.drake_optimization_ik import DrakeOptimizationIK +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus +from dimos.manipulation.planning.spec.models import IKResult +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +class FakeWorld: + def __init__(self) -> None: + self.robot_id = "robot-instance" + self.config = RobotModelConfig( + name="arm", + model_path=Path("/tmp/fake.urdf"), + joint_names=["base", "shoulder", "elbow", "wrist"], + ) + self.current_state = JointState( + {"name": ["base", "shoulder", "elbow", "wrist"], "position": [1.0, 2.0, 3.0, 4.0]} + ) + self.collision_checked_state: JointState | None = None + + def get_robot_ids(self) -> list[str]: + return [self.robot_id] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + assert robot_id == self.robot_id + return self.config + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + assert robot_id == self.robot_id + return np.array([-10.0] * 4), np.array([10.0] * 4) + + @contextmanager + def scratch_context(self): + yield object() + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + assert robot_id == self.robot_id + return self.current_state + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + assert robot_id == self.robot_id + self.collision_checked_state = joint_state + return True + + +def test_solve_pose_targets_uses_group_tip_locks_seed_fallback_and_filters(monkeypatch) -> None: + monkeypatch.setattr(drake_ik, "DRAKE_AVAILABLE", True) + monkeypatch.setattr(DrakeOptimizationIK, "_validate_world", lambda self, world: None) + + world = FakeWorld() + group = PlanningGroup( + id="arm/reach", + robot_name="arm", + group_name="reach", + joint_names=("arm/shoulder", "arm/wrist"), + local_joint_names=("shoulder", "wrist"), + base_link="base_link", + tip_link="group_tip_link", + ) + calls = [] + + def fake_solve_single(self, **kwargs) -> IKResult: + calls.append(kwargs) + return IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + { + "name": ["base", "shoulder", "elbow", "wrist"], + "position": [10.0, 20.0, 30.0, 40.0], + } + ), + position_error=0.0, + orientation_error=0.0, + ) + + monkeypatch.setattr(DrakeOptimizationIK, "_solve_single", fake_solve_single) + monkeypatch.setattr(drake_ik, "RigidTransform", lambda matrix: matrix, raising=False) + + result = DrakeOptimizationIK().solve_pose_targets( + world=world, # type: ignore[arg-type] + pose_targets={group: PoseStamped()}, + seed=JointState({"name": ["shoulder"], "position": [22.0]}), + check_collision=False, + max_attempts=1, + ) + + assert result.is_success() + assert result.joint_state is not None + assert result.joint_state.name == ["arm/shoulder", "arm/wrist"] + assert result.joint_state.position == [20.0, 40.0] + assert calls[0]["target_frame_name"] == "group_tip_link" + np.testing.assert_allclose(calls[0]["seed"], [1.0, 22.0, 3.0, 4.0]) + assert calls[0]["locked_joint_positions"] == {0: 1.0, 2: 3.0} + + +class FakeRigidTransform: + def translation(self) -> np.ndarray: + return np.array([0.5, 0.6, 0.7]) + + def rotation(self) -> str: + return "target-rotation" + + def GetAsMatrix4(self) -> np.ndarray: + return np.eye(4) + + +class FakeBody: + def __init__(self, frame_name: str) -> None: + self.frame_name = frame_name + + def body_frame(self) -> str: + return f"frame:{self.frame_name}" + + +class FakePlant: + def __init__(self) -> None: + self.requested_bodies: list[tuple[str, str]] = [] + + def GetBodyByName(self, frame_name: str, model_instance: str) -> FakeBody: + self.requested_bodies.append((frame_name, model_instance)) + return FakeBody(frame_name) + + def world_frame(self) -> str: + return "world-frame" + + def num_positions(self) -> int: + return 6 + + +class FakeProgram: + def __init__(self) -> None: + self.locks: list[tuple[float, float, str]] = [] + self.initial_guess: tuple[list[str], np.ndarray] | None = None + + def AddBoundingBoxConstraint(self, lower: float, upper: float, variable: str) -> None: + self.locks.append((lower, upper, variable)) + + def SetInitialGuess(self, q: list[str], full_seed: np.ndarray) -> None: + self.initial_guess = (q, full_seed.copy()) + + +class FakeInverseKinematics: + instances: list[FakeInverseKinematics] = [] + + def __init__(self, plant: FakePlant) -> None: + self.plant = plant + self.program = FakeProgram() + self.q_vars = [f"q{i}" for i in range(6)] + self.position_constraints = [] + self.orientation_constraints = [] + self.instances.append(self) + + def AddPositionConstraint(self, **kwargs) -> None: + self.position_constraints.append(kwargs) + + def AddOrientationConstraint(self, **kwargs) -> None: + self.orientation_constraints.append(kwargs) + + def get_mutable_prog(self) -> FakeProgram: + return self.program + + def q(self) -> list[str]: + return self.q_vars + + +class FakeSolveResult: + def is_success(self) -> bool: + return True + + def GetSolution(self, q: list[str]) -> np.ndarray: + return np.array([0.0, 11.0, 0.0, 22.0, 33.0, 0.0]) + + +class FakeDrakeWorld: + def __init__(self) -> None: + self.plant = FakePlant() + self._robots = {"robot-instance": _FakeRobotData()} + self.link_pose_calls: list[tuple[str, str]] = [] + self.set_joint_state_calls: list[JointState] = [] + + @contextmanager + def scratch_context(self): + yield "ctx" + + def set_joint_state(self, ctx: str, robot_id: str, joint_state: JointState) -> None: + self.set_joint_state_calls.append(joint_state) + + def get_link_pose(self, ctx: str, robot_id: str, target_frame_name: str) -> np.ndarray: + self.link_pose_calls.append((robot_id, target_frame_name)) + return np.eye(4) + + +class _FakeRobotData: + model_instance = "model-instance" + joint_indices = [1, 3, 4] + + +def test_solve_single_uses_target_frame_for_constraints_error_and_joint_locks(monkeypatch) -> None: + FakeInverseKinematics.instances.clear() + monkeypatch.setattr(drake_ik, "DRAKE_AVAILABLE", True) + monkeypatch.setattr(drake_ik, "InverseKinematics", FakeInverseKinematics, raising=False) + monkeypatch.setattr(drake_ik, "RotationMatrix", lambda: "identity-rotation", raising=False) + monkeypatch.setattr(drake_ik, "Solve", lambda prog: FakeSolveResult(), raising=False) + monkeypatch.setattr(drake_ik, "compute_pose_error", lambda actual, target: (0.01, 0.02)) + + world = FakeDrakeWorld() + result = DrakeOptimizationIK()._solve_single( + world=world, # type: ignore[arg-type] + robot_id="robot-instance", + target_transform=FakeRigidTransform(), + seed=np.array([1.0, 2.0, 3.0]), + joint_names=["j0", "j1", "j2"], + position_tolerance=0.1, + orientation_tolerance=0.2, + lower_limits=np.array([-5.0, -5.0, -5.0]), + upper_limits=np.array([5.0, 5.0, 5.0]), + target_frame_name="selected_tip_link", + locked_joint_positions={0: 1.5, 2: 3.5}, + ) + + ik = FakeInverseKinematics.instances[0] + assert result.is_success() + assert world.plant.requested_bodies == [("selected_tip_link", "model-instance")] + assert ik.position_constraints[0]["frameB"] == "frame:selected_tip_link" + assert ik.orientation_constraints[0]["frameBbar"] == "frame:selected_tip_link" + assert world.link_pose_calls == [("robot-instance", "selected_tip_link")] + assert ik.program.locks == [(1.5, 1.5, "q1"), (3.5, 3.5, "q4")] + assert ik.program.initial_guess is not None + np.testing.assert_allclose(ik.program.initial_guess[1], [0.0, 1.0, 0.0, 2.0, 3.0, 0.0]) diff --git a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py new file mode 100644 index 0000000000..15d966b7a4 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py @@ -0,0 +1,147 @@ +# 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 tests for group-aware Jacobian IK.""" + +from __future__ import annotations + +from contextlib import nullcontext +from pathlib import Path + +import numpy as np + +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus +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.sensor_msgs.JointState import JointState + + +def _pose(x: float = 0.0) -> PoseStamped: + return PoseStamped(position=Vector3(x, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + + +def _group(tip_link: str | None = "tool") -> PlanningGroup: + return PlanningGroup( + id="arm/manipulator", + robot_name="arm", + group_name="manipulator", + joint_names=("arm/joint_a", "arm/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link=tip_link, + ) + + +class _World: + is_finalized = True + + def __init__(self) -> None: + self.group_pose_calls = 0 + self.group_jacobian_calls = 0 + self.legacy_pose_calls = 0 + self.legacy_jacobian_calls = 0 + self.config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + base_pose=_pose(), + joint_names=["joint_a", "joint_b", "gripper"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ) + ], + ) + + def get_robot_ids(self) -> list[str]: + return ["robot"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.config + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "gripper"], "position": [0.0, 0.0, 0.9]}) + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.last_state = joint_state + + def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: + self.group_pose_calls += 1 + return _pose(0.0) + + def get_group_jacobian(self, ctx: object, group_id: str) -> np.ndarray: + self.group_jacobian_calls += 1 + return np.eye(6, 2) + + def get_ee_pose(self, ctx: object, robot_id: str) -> PoseStamped: + self.legacy_pose_calls += 1 + raise AssertionError("legacy EE pose should not be used") + + def get_jacobian(self, ctx: object, robot_id: str) -> np.ndarray: + self.legacy_jacobian_calls += 1 + raise AssertionError("legacy Jacobian should not be used") + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + return True + + +def test_solve_pose_targets_filters_to_group_and_uses_group_world_methods() -> None: + world = _World() + result = JacobianIK(max_iterations=2).solve_pose_targets( + world=world, + pose_targets={_group(): _pose()}, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": [0.0, 0.0, 0.9]} + ), + max_attempts=1, + ) + + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b"] + assert world.group_pose_calls == 1 + assert world.group_jacobian_calls == 0 + assert world.legacy_pose_calls == 0 + assert world.legacy_jacobian_calls == 0 + + +def test_solve_pose_targets_rejects_auxiliary_groups() -> None: + result = JacobianIK().solve_pose_targets( + world=_World(), + pose_targets={_group(): _pose()}, + auxiliary_groups=[_group()], + ) + + assert result.status == IKStatus.UNSUPPORTED + assert "no auxiliary" in result.message + + +def test_solve_pose_targets_rejects_group_without_pose_target_frame() -> None: + result = JacobianIK().solve_pose_targets(world=_World(), pose_targets={_group(None): _pose()}) + + assert result.status == IKStatus.UNSUPPORTED + assert "no pose target frame" in result.message diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index 920f408e22..d0324edb56 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -26,7 +26,7 @@ from pytest_mock import MockerFixture from dimos.manipulation.planning.factory import create_kinematics -from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig import dimos.manipulation.planning.kinematics.pink_ik as pink_ik from dimos.manipulation.planning.kinematics.pink_ik import ( @@ -54,8 +54,9 @@ def __init__(self, idx_q: int) -> None: class _FakeFrame: - def __init__(self, name: str) -> None: + def __init__(self, name: str, parent_joint: int = 0) -> None: self.name = name + self.parentJoint = parent_joint class _FakePlacement: @@ -67,7 +68,7 @@ def __init__(self, translation: np.ndarray) -> None: class _FakeData: def __init__(self) -> None: self.q = np.zeros(3) - self.oMf = [_FakePlacement(np.zeros(3))] + self.oMf = [_FakePlacement(np.zeros(3)), _FakePlacement(np.zeros(3))] class _FakeModel: @@ -76,9 +77,9 @@ class _FakeModel: def __init__(self) -> None: self.names = ["universe", "joint_b", "joint_a", "joint_c"] self.joints = [SimpleNamespace(idx_q=-1, nq=0), _FakeJoint(0), _FakeJoint(1), _FakeJoint(2)] - self.frames = [_FakeFrame("tool")] + self.frames = [_FakeFrame("base", 0), _FakeFrame("tool", 3)] self._joint_ids = {"joint_b": 1, "joint_a": 2, "joint_c": 3} - self._frame_ids = {"tool": 0} + self._frame_ids = {"base": 0, "tool": 1} def createData(self) -> _FakeData: return _FakeData() @@ -138,7 +139,7 @@ def forward_kinematics(model: _FakeModel, data: _FakeData, q: np.ndarray) -> Non data.q = q.copy() def update_frame_placements(model: _FakeModel, data: _FakeData) -> None: - data.oMf[0] = _FakePlacement(data.q.copy()) + data.oMf[1] = _FakePlacement(data.q.copy()) pinocchio.forwardKinematics = forward_kinematics # type: ignore[attr-defined] pinocchio.updateFramePlacements = update_frame_placements # type: ignore[attr-defined] @@ -195,7 +196,7 @@ def _context() -> _PinkRobotContext: return _PinkRobotContext( model=model, data=model.createData(), - frame_id=0, + frame_id=1, frame_name="tool", mapping=mapping, ) @@ -207,6 +208,39 @@ class _FakeWorld: def __init__(self, collision_free: bool = True) -> None: self.config = _robot_config() self.collision_free = collision_free + self.joint_state_calls = 0 + self.groups = { + "arm/manipulator": PlanningGroup( + id="arm/manipulator", + robot_name="arm", + group_name="manipulator", + joint_names=("arm/joint_a", "arm/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + "arm/no_tip": PlanningGroup( + id="arm/no_tip", + robot_name="arm", + group_name="no_tip", + joint_names=("arm/joint_c",), + local_joint_names=("joint_c",), + base_link="base", + tip_link=None, + ), + "arm/wrist": PlanningGroup( + id="arm/wrist", + robot_name="arm", + group_name="wrist", + joint_names=("arm/joint_c",), + local_joint_names=("joint_c",), + base_link="base", + tip_link="base", + ), + } + + def get_robot_ids(self) -> list[str]: + return ["robot"] def get_robot_config(self, robot_id: str) -> RobotModelConfig: return self.config @@ -215,10 +249,8 @@ def scratch_context(self) -> nullcontext[None]: return nullcontext(None) def get_joint_state(self, ctx: object, robot_id: str) -> JointState: - return JointState( - name=["joint_b", "joint_c", "joint_a"], - position=[0.0, 0.0, 0.0], - ) + self.joint_state_calls += 1 + return JointState({"name": ["joint_b", "joint_c", "joint_a"], "position": [0.0, 0.0, 0.0]}) def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -226,6 +258,72 @@ def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: return self.collision_free + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.joint_state = joint_state + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + return self.collision_free + + +class _MultiRobotCollisionWorld: + is_finalized = True + + def __init__(self) -> None: + left_config = _robot_config() + left_config.name = "left" + right_config = _robot_config() + right_config.name = "right" + self.configs = {"left-id": left_config, "right-id": right_config} + self.groups = { + "left/manipulator": PlanningGroup( + id="left/manipulator", + robot_name="left", + group_name="manipulator", + joint_names=("left/joint_a", "left/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + "right/manipulator": PlanningGroup( + id="right/manipulator", + robot_name="right", + group_name="manipulator", + joint_names=("right/joint_a", "right/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + } + self.config_collision_checks = 0 + self.context_collision_checks = 0 + self.context_states: dict[str, JointState] = {} + + def get_robot_ids(self) -> list[str]: + return ["left-id", "right-id"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.configs[robot_id] + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}) + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + self.config_collision_checks += 1 + return True + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.context_states[robot_id] = joint_state + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + self.context_collision_checks += 1 + return len(self.context_states) < 2 + def test_create_kinematics_pink_missing_dependency_is_actionable( mocker: MockerFixture, @@ -349,7 +447,7 @@ def test_solve_single_reports_non_convergence(mocker: MockerFixture) -> None: def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: ik = _pink_ik(mocker, converge=True) context = _context() - ik._robot_contexts = {"robot": context} + ik._robot_contexts = {("robot", "tool"): context} result = ik.solve( world=cast("Any", _FakeWorld(collision_free=False)), @@ -369,7 +467,7 @@ def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: def test_solve_retries_after_joint_limit_failure(mocker: MockerFixture) -> None: ik = _pink_ik(mocker, converge=True) context = _context() - ik._robot_contexts = {"robot": context} + ik._robot_contexts = {("robot", "tool"): context} calls = 0 def fake_solve_single(**_: object) -> IKResult: @@ -407,3 +505,245 @@ def fake_solve_single(**_: object) -> IKResult: assert solve_single.call_count == 2 assert result.status == IKStatus.SUCCESS + + +def test_robot_context_cache_key_includes_tip_frame(mocker: MockerFixture, tmp_path: Path) -> None: + modules = _fake_modules() + modules.pinocchio.buildModelFromUrdf = lambda path: _FakeModel() # type: ignore[attr-defined] + mocker.patch.object(pink_ik, "_load_optional_dependencies", return_value=modules) + mocker.patch.object(pink_ik, "prepare_urdf_for_drake", return_value=tmp_path / "prepared.urdf") + model_path = tmp_path / "fake.urdf" + model_path.write_text("") + world = _FakeWorld() + world.config.model_path = model_path + ik = PinkIK(PinkIKConfig(max_iterations=1)) + + first = ik._get_robot_context(cast("Any", world), "robot", "tool") + second = ik._get_robot_context(cast("Any", world), "robot", "base") + + assert first is not second + assert set(ik._robot_contexts) == {("robot", "tool"), ("robot", "base")} + + +def test_build_robot_context_rejects_base_link_not_model_root( + mocker: MockerFixture, tmp_path: Path +) -> None: + model = _FakeModel() + model.frames[0] = _FakeFrame("base", parent_joint=1) + modules = _fake_modules() + modules.pinocchio.buildModelFromUrdf = lambda path: model # type: ignore[attr-defined] + mocker.patch.object(pink_ik, "_load_optional_dependencies", return_value=modules) + mocker.patch.object(pink_ik, "prepare_urdf_for_drake", return_value=tmp_path / "prepared.urdf") + model_path = tmp_path / "fake.urdf" + model_path.write_text("") + config = _robot_config() + config.model_path = model_path + + with pytest.raises(ValueError, match="base_link 'base'.*model root"): + PinkIK(PinkIKConfig(max_iterations=1))._build_robot_context(config, "tool") + + +def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker, converge=True) + context = _context() + get_context = mocker.patch.object(ik, "_get_robot_context", return_value=context) + mocker.patch.object( + ik, + "_solve_single", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + ) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.0, 0.0, 0.0]} + ), + max_attempts=1, + ) + + get_context.assert_called_once_with(cast("Any", world), "robot", "tool") + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b"] + assert result.joint_state.position == [0.1, 0.2] + assert world.joint_state_calls == 0 + + +def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/no_tip"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + ) + + assert result.status == IKStatus.UNSUPPORTED + assert "no pose target frame" in result.message + + +def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + mocker.patch.object( + ik, + "_solve_single", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + ) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + seed=JointState({"name": ["arm/joint_a"], "position": [0.0]}), + max_attempts=1, + ) + + assert result.status == IKStatus.SUCCESS + assert world.joint_state_calls == 1 + + +def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + solve_multi = mocker.patch.object( + ik, + "_solve_multi", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + position_error=0.0, + orientation_error=0.0, + ), + ) + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + world.groups["arm/wrist"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + }, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.0, 0.0, 0.0]} + ), + max_attempts=1, + ) + + solve_multi.assert_called_once() + assert len(solve_multi.call_args.kwargs["targets"]) == 2 + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b", "arm/joint_c"] + assert result.joint_state.position == [0.1, 0.2, 0.3] + + +def test_solve_pose_targets_checks_multi_robot_solution_together( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker) + world = _MultiRobotCollisionWorld() + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + solve_single = mocker.patch.object( + ik, + "_solve_single", + side_effect=[ + IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.4, 0.5, 0.6]} + ), + ), + ], + ) + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["left/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + world.groups["right/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + }, + seed=JointState( + { + "name": [ + "left/joint_a", + "left/joint_b", + "left/joint_c", + "right/joint_a", + "right/joint_b", + "right/joint_c", + ], + "position": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + } + ), + max_attempts=1, + ) + + assert solve_single.call_count == 2 + assert result.status == IKStatus.COLLISION + assert world.config_collision_checks == 0 + assert world.context_collision_checks == 1 + assert set(world.context_states) == {"left-id", "right-id"} + + +def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={}, + auxiliary_groups=[world.groups["arm/no_tip"], world.groups["arm/manipulator"]], + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ) + + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_c", "arm/joint_a", "arm/joint_b"] + assert result.joint_state.position == [0.3, 0.1, 0.2] + assert world.joint_state_calls == 0 diff --git a/dimos/manipulation/planning/kinematics/utils.py b/dimos/manipulation/planning/kinematics/utils.py new file mode 100644 index 0000000000..3d0fa0509c --- /dev/null +++ b/dimos/manipulation/planning/kinematics/utils.py @@ -0,0 +1,236 @@ +# 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. + +"""Shared IK-only helpers for planning-group-scoped kinematics backends.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints +from dimos.manipulation.planning.spec.enums import IKStatus +from dimos.manipulation.planning.spec.models import IKResult, RobotName, WorldRobotID +from dimos.manipulation.planning.spec.protocols import WorldSpec +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +@dataclass(frozen=True) +class SinglePoseTargetRequest: + group: PlanningGroup + target_pose: PoseStamped + robot_id: WorldRobotID + joint_names: list[str] + seed_positions: NDArray[np.float64] + group_indices: list[int] + + +def unique_pose_target_frame_for_robot(world: WorldSpec, robot_id: WorldRobotID) -> str | None: + config = world.get_robot_config(robot_id) + pose_target_frames = [ + group.tip_link for group in config.planning_groups if group.tip_link is not None + ] + unique_frames = list(dict.fromkeys(pose_target_frames)) + if len(unique_frames) != 1: + return None + return unique_frames[0] + + +def robot_ids_by_name( + world: WorldSpec, + robot_names: tuple[RobotName, ...], +) -> dict[RobotName, WorldRobotID]: + robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + for robot_name in robot_names: + matches = [ + robot_id + for robot_id in world.get_robot_ids() + if world.get_robot_config(robot_id).name == robot_name + ] + if not matches: + raise ValueError(f"Robot '{robot_name}' not found") + if len(matches) > 1: + raise ValueError(f"Robot name '{robot_name}' is not unique in planning world") + robot_ids_by_name[robot_name] = matches[0] + return robot_ids_by_name + + +def seed_positions_with_world_fallback( + world: WorldSpec, + robot_id: WorldRobotID, + robot_name: RobotName, + local_joint_names: list[str], + seed: JointState | None, +) -> NDArray[np.float64]: + """Return full robot positions, reading world only for absent seed joints.""" + if seed is None: + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx, robot_id) + return positions_by_local_name(current, robot_name, local_joint_names) + + try: + return positions_by_local_name(seed, robot_name, local_joint_names) + except ValueError: + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx, robot_id) + fallback_positions = positions_by_local_name(current, robot_name, local_joint_names) + seed_positions = partial_positions_by_local_name(seed, robot_name, local_joint_names) + local_indices = {name: index for index, name in enumerate(local_joint_names)} + for local_name, position in seed_positions.items(): + fallback_positions[local_indices[local_name]] = position + return fallback_positions + + +def resolve_single_pose_target_request( + world: WorldSpec, + pose_targets: dict[PlanningGroup, PoseStamped] | Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup], + seed: JointState | None, + backend_name: str, +) -> tuple[SinglePoseTargetRequest | None, IKResult | None]: + if not pose_targets: + return None, _failure(IKStatus.NO_SOLUTION, "At least one pose target is required") + if len(pose_targets) != 1 or auxiliary_groups: + return None, _failure( + IKStatus.UNSUPPORTED, + f"{backend_name} supports exactly one pose target and no auxiliary planning groups", + ) + + target_group = next(iter(pose_targets.keys())) + if not target_group.has_pose_target: + return None, _failure( + IKStatus.UNSUPPORTED, + f"Planning group '{target_group.id}' has no pose target frame", + ) + + try: + selection = PlanningGroupSelection.from_groups((target_group,)) + robot_id = robot_ids_by_name(world, selection.robot_names)[target_group.robot_name] + config = world.get_robot_config(robot_id) + joint_names = list(config.joint_names) + seed_positions = seed_positions_with_world_fallback( + world, + robot_id, + config.name, + joint_names, + seed, + ) + group_indices = [joint_names.index(name) for name in target_group.local_joint_names] + except ValueError as exc: + return None, _failure(IKStatus.NO_SOLUTION, str(exc)) + + return ( + SinglePoseTargetRequest( + group=target_group, + target_pose=pose_targets[target_group], + robot_id=robot_id, + joint_names=joint_names, + seed_positions=seed_positions, + group_indices=group_indices, + ), + None, + ) + + +def positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[str], +) -> NDArray[np.float64]: + if not joint_state.name: + if len(joint_state.position) != len(local_joint_names): + raise ValueError( + f"JointState has {len(joint_state.position)} positions for " + f"{len(local_joint_names)} joints" + ) + return np.asarray(joint_state.position, dtype=np.float64) + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + positions: list[float] = [] + missing: list[str] = [] + for local_name in local_joint_names: + global_name = f"{robot_name}/{local_name}" + if local_name in positions_by_name: + positions.append(float(positions_by_name[local_name])) + elif global_name in positions_by_name: + positions.append(float(positions_by_name[global_name])) + else: + missing.append(local_name) + if missing: + raise ValueError(f"JointState missing joints: {missing}") + return np.asarray(positions, dtype=np.float64) + + +def partial_positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[str], +) -> dict[str, float]: + if len(joint_state.name) != len(joint_state.position): + raise ValueError( + f"Seed has {len(joint_state.name)} names but {len(joint_state.position)} positions" + ) + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + known_local_names = set(local_joint_names) + positions: dict[str, float] = {} + for name, position in positions_by_name.items(): + if name in known_local_names: + positions[name] = float(position) + continue + prefix = f"{robot_name}/" + if name.startswith(prefix): + local_name = name[len(prefix) :] + if local_name in known_local_names: + positions[local_name] = float(position) + continue + raise ValueError(f"Unrecognized seed joint '{name}'") + return positions + + +def filter_result_to_group(result: IKResult, group: PlanningGroup) -> IKResult: + return filter_result_to_selection(result, PlanningGroupSelection.from_groups((group,))) + + +def filter_result_to_selection(result: IKResult, selection: PlanningGroupSelection) -> IKResult: + if result.joint_state is None: + return result + local_joint_names = tuple( + local_name for group in selection.groups for local_name in group.local_joint_names + ) + return IKResult( + status=result.status, + joint_state=filter_joint_state_to_selected_joints( + result.joint_state, + selection.joint_names, + local_joint_names, + ), + position_error=result.position_error, + orientation_error=result.orientation_error, + iterations=result.iterations, + message=result.message, + ) + + +def groups_by_robot(groups: Sequence[PlanningGroup]) -> dict[RobotName, list[PlanningGroup]]: + grouped: dict[RobotName, list[PlanningGroup]] = {} + for group in groups: + grouped.setdefault(group.robot_name, []).append(group) + return grouped + + +def _failure(status: IKStatus, message: str) -> IKResult: + return IKResult(status=status, joint_state=None, message=message) diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 3ca19eb099..a98ea358e6 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -26,8 +26,17 @@ import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroupSelection +from dimos.manipulation.planning.planners.selected_joint_space import ( + SelectedJointSpace, + normalize_selection_target, +) from dimos.manipulation.planning.spec.enums import PlanningStatus -from dimos.manipulation.planning.spec.models import JointPath, PlanningResult, WorldRobotID +from dimos.manipulation.planning.spec.models import ( + JointPath, + PlanningResult, + WorldRobotID, +) from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.msgs.sensor_msgs.JointState import JointState @@ -98,6 +107,9 @@ def plan_joint_path( if error is not None: return error + if world.check_edge_collision_free(robot_id, start, goal, self._collision_step_size): + return _create_success_result([start, goal], time.time() - start_time, 0) + lower, upper = world.get_joint_limits(robot_id) start_tree = [TreeNode(config=q_start.copy())] goal_tree = [TreeNode(config=q_goal.copy())] @@ -147,6 +159,183 @@ def get_name(self) -> str: """Get planner name.""" return "RRTConnect" + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan over an explicit planning-group selection. + + The search space is the selected global-joint order. Collision checks project + candidates into full per-robot states, holding unselected joints at the world + current state. + """ + start_time = time.time() + if not world.is_finalized: + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + "World must be finalized before planning", + ) + + if not selection.groups: + return _create_failure_result( + PlanningStatus.INVALID_GOAL, "No planning groups selected" + ) + + selected_joint_names = list(selection.joint_names) + try: + normalized_start = normalize_selection_target(selection, start, "start") + except ValueError as exc: + return _create_failure_result(PlanningStatus.INVALID_START, str(exc)) + try: + normalized_goal = normalize_selection_target(selection, goal, "goal") + except ValueError as exc: + return _create_failure_result(PlanningStatus.INVALID_GOAL, str(exc)) + try: + selected_space = SelectedJointSpace.from_world(world, selection) + q_start = np.asarray(normalized_start.position, dtype=np.float64) + q_goal = np.asarray(normalized_goal.position, dtype=np.float64) + lower, upper = selected_space.joint_limits() + except ValueError as exc: + return _create_failure_result(PlanningStatus.NO_SOLUTION, str(exc)) + + if np.any(q_start < lower) or np.any(q_start > upper): + return _create_failure_result( + PlanningStatus.INVALID_START, + "Start configuration is outside joint limits", + ) + if np.any(q_goal < lower) or np.any(q_goal > upper): + return _create_failure_result( + PlanningStatus.INVALID_GOAL, + "Goal configuration is outside joint limits", + ) + + if not selected_space.config_collision_free(world, q_start): + return _create_failure_result( + PlanningStatus.COLLISION_AT_START, + "Start configuration is in collision", + ) + if not selected_space.config_collision_free(world, q_goal): + return _create_failure_result( + PlanningStatus.COLLISION_AT_GOAL, + "Goal configuration is in collision", + ) + + if selected_space.edge_collision_free( + world, + q_start, + q_goal, + self._collision_step_size, + ): + return _create_success_result( + [normalized_start, normalized_goal], time.time() - start_time, 0 + ) + + start_tree = [TreeNode(config=q_start.copy())] + goal_tree = [TreeNode(config=q_goal.copy())] + trees_swapped = False + + for iteration in range(max_iterations): + if time.time() - start_time > timeout: + return _create_failure_result( + PlanningStatus.TIMEOUT, + f"Timeout after {iteration} iterations", + time.time() - start_time, + iteration, + ) + + sample = np.random.uniform(lower, upper) + extended = self._extend_selected_tree( + selected_space, + world, + start_tree, + sample, + self._step_size, + ) + if extended is not None: + connected = self._connect_selected_tree( + selected_space, + world, + goal_tree, + extended.config, + self._connect_step_size, + ) + if connected is not None: + path = self._extract_path(extended, connected, selected_joint_names) + if trees_swapped: + path = list(reversed(path)) + path = selected_space.simplify_path( + world, + path, + self._collision_step_size, + ) + return _create_success_result(path, time.time() - start_time, iteration + 1) + + start_tree, goal_tree = goal_tree, start_tree + trees_swapped = not trees_swapped + + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + f"No path found after {max_iterations} iterations", + time.time() - start_time, + max_iterations, + ) + + def _extend_selected_tree( + self, + selected_space: SelectedJointSpace, + world: WorldSpec, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Extend a tree in selected-joint space.""" + nearest = min(tree, key=lambda node: float(np.linalg.norm(node.config - target))) + diff = target - nearest.config + dist = float(np.linalg.norm(diff)) + if dist <= step_size: + new_config = target.copy() + else: + new_config = nearest.config + step_size * (diff / dist) + + if selected_space.edge_collision_free( + world, + nearest.config, + new_config, + self._collision_step_size, + ): + new_node = TreeNode(config=new_config, parent=nearest) + nearest.children.append(new_node) + tree.append(new_node) + return new_node + return None + + def _connect_selected_tree( + self, + selected_space: SelectedJointSpace, + world: WorldSpec, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Try to connect a selected-joint tree to a target.""" + while True: + result = self._extend_selected_tree( + selected_space, + world, + tree, + target, + step_size, + ) + if result is None: + return None + if float(np.linalg.norm(result.config - target)) < self._goal_tolerance: + return result + def _validate_inputs( self, world: WorldSpec, @@ -226,8 +415,8 @@ def _extend_tree( new_config = nearest.config + step_size * (diff / dist) # Check validity of edge using context-free method - start_state = JointState(name=joint_names, position=nearest.config.tolist()) - end_state = JointState(name=joint_names, position=new_config.tolist()) + start_state = JointState({"name": joint_names, "position": nearest.config.tolist()}) + end_state = JointState({"name": joint_names, "position": new_config.tolist()}) if world.check_edge_collision_free( robot_id, start_state, end_state, self._collision_step_size ): @@ -277,7 +466,7 @@ def _extract_path( full_path_arrays = start_path + list(reversed(goal_path)) # Convert to list of JointState - return [JointState(name=joint_names, position=q.tolist()) for q in full_path_arrays] + return [JointState({"name": joint_names, "position": q.tolist()}) for q in full_path_arrays] def _simplify_path( self, diff --git a/dimos/manipulation/planning/planners/selected_joint_space.py b/dimos/manipulation/planning/planners/selected_joint_space.py new file mode 100644 index 0000000000..89d6efbadc --- /dev/null +++ b/dimos/manipulation/planning/planners/selected_joint_space.py @@ -0,0 +1,319 @@ +# 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. + +"""Selected planning-group joint-space projection helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + local_joint_name_from_global, + make_global_joint_name, +) +from dimos.manipulation.planning.groups.models import PlanningGroupSelection +from dimos.manipulation.planning.spec.models import ( + JointPath, + LocalModelJointName, + RobotName, + WorldRobotID, +) +from dimos.manipulation.planning.spec.protocols import WorldSpec +from dimos.msgs.sensor_msgs.JointState import JointState + + +@dataclass(frozen=True) +class SelectedRobotProjection: + """Runtime state needed to project selected-group samples into one robot. + + This is not static model metadata like ``RobotModelConfig``. It captures the + current planning context: the world robot ID, base/current positions used for + non-selected joints, and joint-limit lookup tables for the selected-space + planner. + """ + + robot_id: WorldRobotID + robot_name: RobotName + local_joint_names: list[LocalModelJointName] + base_positions_by_local_name: dict[LocalModelJointName, float] + lower_limits_by_local_name: dict[LocalModelJointName, float] + upper_limits_by_local_name: dict[LocalModelJointName, float] + + +class SelectedJointSpace: + """Projection adapter between selected global joints and full robot states.""" + + def __init__( + self, + robot_projections: list[SelectedRobotProjection], + selected_joint_names: list[str], + ) -> None: + self.robot_projections = robot_projections + self.selected_joint_names = selected_joint_names + + @classmethod + def from_world( + cls, + world: WorldSpec, + selection: PlanningGroupSelection, + ) -> SelectedJointSpace: + return cls( + robot_projections=_build_robot_projections(world, selection), + selected_joint_names=list(selection.joint_names), + ) + + def joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + projections_by_robot_name = { + projection.robot_name: projection for projection in self.robot_projections + } + lower: list[float] = [] + upper: list[float] = [] + for global_name in self.selected_joint_names: + robot_name, local_name = _split_selected_global_joint_name( + global_name, projections_by_robot_name + ) + projection = projections_by_robot_name[robot_name] + lower.append(projection.lower_limits_by_local_name[local_name]) + upper.append(projection.upper_limits_by_local_name[local_name]) + return np.asarray(lower, dtype=np.float64), np.asarray(upper, dtype=np.float64) + + def config_collision_free( + self, + world: WorldSpec, + selected_positions: NDArray[np.float64], + ) -> bool: + with world.scratch_context() as ctx: + projected_states = self.project_config(selected_positions) + for projection in self.robot_projections: + world.set_joint_state( + ctx, projection.robot_id, projected_states[projection.robot_id] + ) + return all( + world.is_collision_free(ctx, projection.robot_id) + for projection in self.robot_projections + ) + + def edge_collision_free( + self, + world: WorldSpec, + start: NDArray[np.float64], + end: NDArray[np.float64], + step_size: float, + ) -> bool: + distance = float(np.linalg.norm(end - start)) + steps = max(1, int(np.ceil(distance / step_size))) + for step in range(steps + 1): + ratio = step / steps + candidate = start + ratio * (end - start) + if not self.config_collision_free(world, candidate): + return False + return True + + def project_config( + self, + selected_positions: NDArray[np.float64], + ) -> dict[WorldRobotID, JointState]: + selected_positions_by_global_name = dict( + zip(self.selected_joint_names, selected_positions.tolist(), strict=True) + ) + projected_states: dict[WorldRobotID, JointState] = {} + for projection in self.robot_projections: + positions: list[float] = [] + for local_name in projection.local_joint_names: + global_name = make_global_joint_name(projection.robot_name, local_name) + position = selected_positions_by_global_name.get( + global_name, + projection.base_positions_by_local_name[local_name], + ) + positions.append(float(position)) + projected_states[projection.robot_id] = JointState( + {"name": list(projection.local_joint_names), "position": positions} + ) + return projected_states + + def simplify_path( + self, + world: WorldSpec, + path: JointPath, + collision_step_size: float, + max_iterations: int = 100, + ) -> JointPath: + if len(path) <= 2: + return path + + simplified = list(path) + for _ in range(max_iterations): + if len(simplified) <= 2: + break + i = np.random.randint(0, len(simplified) - 2) + j = np.random.randint(i + 2, len(simplified)) + start = np.asarray(simplified[i].position, dtype=np.float64) + end = np.asarray(simplified[j].position, dtype=np.float64) + if self.edge_collision_free(world, start, end, collision_step_size): + simplified = simplified[: i + 1] + simplified[j:] + return simplified + + +def normalize_selection_target( + selection: PlanningGroupSelection, + target: JointState, + label: str, +) -> JointState: + """Normalize a selected-joint target to global selection order.""" + selected_global_names = list(selection.joint_names) + if not target.name: + if len(target.position) != len(selected_global_names): + raise ValueError( + f"{label} target has {len(target.position)} positions, " + f"expected {len(selected_global_names)}" + ) + return JointState({"name": selected_global_names, "position": list(target.position)}) + + if len(target.name) != len(target.position): + raise ValueError( + f"{label} target has {len(target.name)} names but {len(target.position)} positions" + ) + + names = list(target.name) + global_flags = [is_global_joint_name(name) for name in names] + if any(global_flags) and not all(global_flags): + raise ValueError(f"{label} target mixes global and local joint names: {names}") + + if all(global_flags): + expected_names = selected_global_names + else: + if len(selection.groups) != 1: + raise ValueError( + f"{label} target uses local joint names for a multi-group selection; " + "use global joint names" + ) + expected_names = list(selection.groups[0].local_joint_names) + + positions_by_name = dict(zip(names, target.position, strict=True)) + missing = [name for name in expected_names if name not in positions_by_name] + if missing: + raise ValueError(f"{label} target is missing joints: {missing}") + extra = sorted(set(names) - set(expected_names)) + if extra: + raise ValueError(f"{label} target has extra joints: {extra}") + + ordered_positions = [float(positions_by_name[name]) for name in expected_names] + return JointState({"name": selected_global_names, "position": ordered_positions}) + + +def _build_robot_projections( + world: WorldSpec, + selection: PlanningGroupSelection, +) -> list[SelectedRobotProjection]: + robot_ids_by_name = _robot_ids_by_name(world, selection.robot_names) + robot_projections: list[SelectedRobotProjection] = [] + with world.scratch_context() as ctx: + for robot_name in selection.robot_names: + robot_id = robot_ids_by_name[robot_name] + config = world.get_robot_config(robot_id) + local_joint_names = list(config.joint_names) + current_state = world.get_joint_state(ctx, robot_id) + base_positions_by_local_name = _positions_by_local_name( + current_state, + robot_name, + local_joint_names, + ) + lower, upper = world.get_joint_limits(robot_id) + if len(lower) != len(local_joint_names) or len(upper) != len(local_joint_names): + raise ValueError( + f"Robot '{robot_name}' joint limits do not match configured joints" + ) + robot_projections.append( + SelectedRobotProjection( + robot_id=robot_id, + robot_name=robot_name, + local_joint_names=local_joint_names, + base_positions_by_local_name=base_positions_by_local_name, + lower_limits_by_local_name=dict( + zip(local_joint_names, lower.tolist(), strict=True) + ), + upper_limits_by_local_name=dict( + zip(local_joint_names, upper.tolist(), strict=True) + ), + ) + ) + return robot_projections + + +def _positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[LocalModelJointName], +) -> dict[LocalModelJointName, float]: + if not joint_state.name: + if len(joint_state.position) != len(local_joint_names): + raise ValueError( + f"Current state for robot '{robot_name}' has {len(joint_state.position)} positions, " + f"expected {len(local_joint_names)}" + ) + return dict(zip(local_joint_names, map(float, joint_state.position), strict=True)) + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + positions_by_local_name: dict[LocalModelJointName, float] = {} + for local_name in local_joint_names: + global_name = make_global_joint_name(robot_name, local_name) + if local_name in positions_by_name: + positions_by_local_name[local_name] = float(positions_by_name[local_name]) + elif global_name in positions_by_name: + positions_by_local_name[local_name] = float(positions_by_name[global_name]) + else: + raise ValueError( + f"Current state for robot '{robot_name}' is missing joint '{local_name}'" + ) + return positions_by_local_name + + +def _robot_ids_by_name( + world: WorldSpec, + robot_names: tuple[RobotName, ...], +) -> dict[RobotName, WorldRobotID]: + robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + for robot_name in robot_names: + matches = [ + robot_id + for robot_id in world.get_robot_ids() + if world.get_robot_config(robot_id).name == robot_name + ] + if not matches: + raise ValueError(f"Robot '{robot_name}' not found") + if len(matches) > 1: + raise ValueError(f"Robot name '{robot_name}' is not unique in planning world") + robot_ids_by_name[robot_name] = matches[0] + return robot_ids_by_name + + +def _split_selected_global_joint_name( + global_name: str, + projections_by_robot_name: dict[RobotName, SelectedRobotProjection], +) -> tuple[RobotName, LocalModelJointName]: + for robot_name, projection in projections_by_robot_name.items(): + try: + local_name = local_joint_name_from_global(robot_name, global_name) + except ValueError: + continue + if local_name not in projection.local_joint_names: + raise ValueError( + f"Selected joint '{global_name}' is not configured for robot '{robot_name}'" + ) + return robot_name, local_name + raise ValueError(f"Selected joint '{global_name}' does not belong to a selected robot") diff --git a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py new file mode 100644 index 0000000000..4b330bc816 --- /dev/null +++ b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py @@ -0,0 +1,199 @@ +# 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 tests for selected-joint RRT planning.""" + +from __future__ import annotations + +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import pytest + +from dimos.manipulation.planning.groups.models import ( + PlanningGroup, + PlanningGroupDefinition, + PlanningGroupSelection, +) +from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +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.sensor_msgs.JointState import JointState + + +def _pose() -> PoseStamped: + return PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + + +def _group(name: str, joints: tuple[str, ...]) -> PlanningGroup: + return PlanningGroup( + id=f"arm/{name}", + robot_name="arm", + group_name=name, + joint_names=tuple(f"arm/{joint}" for joint in joints), + local_joint_names=joints, + base_link="base", + tip_link="tool", + ) + + +class _World: + is_finalized = True + + def __init__(self, current: list[float] | None = None) -> None: + self.current = current or [0.0, 0.0, 0.7] + self.projected_states: list[JointState] = [] + self.config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + base_pose=_pose(), + joint_names=["joint_a", "joint_b", "gripper"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition("arm", ("joint_a", "joint_b"), "base", "tool") + ], + ) + + def get_robot_ids(self) -> list[str]: + return ["robot"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.config + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "gripper"], "position": self.current}) + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.projected_states.append(joint_state) + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + return True + + +@pytest.mark.parametrize( + ("start", "goal", "expected_start", "expected_goal"), + [ + ( + JointState({"position": [0.1, 0.2]}), + JointState({"position": [0.3, 0.4]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ( + JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.2, 0.1]}), + JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.4, 0.3]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ( + JointState({"name": ["joint_b", "joint_a"], "position": [0.2, 0.1]}), + JointState({"name": ["joint_b", "joint_a"], "position": [0.4, 0.3]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ], +) +def test_plan_selected_joint_path_normalizes_target_forms( + start: JointState, goal: JointState, expected_start: list[float], expected_goal: list[float] +) -> None: + group = _group("arm", ("joint_a", "joint_b")) + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), PlanningGroupSelection.from_groups((group,)), start, goal + ) + + assert result.status == PlanningStatus.SUCCESS + assert result.path is not None + assert result.path[0].name == ["arm/joint_a", "arm/joint_b"] + assert result.path[0].position == expected_start + assert result.path[-1].position == expected_goal + + +@pytest.mark.parametrize( + ("start", "goal", "status", "message"), + [ + ( + JointState({"name": ["arm/joint_a"], "position": [0.0]}), + JointState({"position": [0.0, 0.0]}), + PlanningStatus.INVALID_START, + "missing", + ), + ( + JointState({"position": [0.0, 0.0]}), + JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/extra"], "position": [0.0, 0.0, 0.0]} + ), + PlanningStatus.INVALID_GOAL, + "extra", + ), + ( + JointState({"name": ["arm/joint_a", "joint_b"], "position": [0.0, 0.0]}), + JointState({"position": [0.0, 0.0]}), + PlanningStatus.INVALID_START, + "mixes", + ), + ], +) +def test_plan_selected_joint_path_rejects_bad_targets( + start: JointState, goal: JointState, status: PlanningStatus, message: str +) -> None: + group = _group("arm", ("joint_a", "joint_b")) + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), PlanningGroupSelection.from_groups((group,)), start, goal + ) + + assert result.status == status + assert message in result.message + + +def test_plan_selected_joint_path_rejects_local_names_for_multi_group_selection() -> None: + selection = PlanningGroupSelection.from_groups( + (_group("arm", ("joint_a",)), _group("gripper", ("gripper",))) + ) + + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), + selection, + JointState({"name": ["joint_a", "gripper"], "position": [0.0, 0.0]}), + JointState({"position": [0.1, 0.2]}), + ) + + assert result.status == PlanningStatus.INVALID_START + assert "multi-group" in result.message + + +def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselected_joints() -> None: + world = _World(current=[0.0, 0.0, 0.77]) + group = _group("arm", ("joint_a", "joint_b")) + + result = RRTConnectPlanner().plan_selected_joint_path( + world, + PlanningGroupSelection.from_groups((group,)), + JointState({"position": [0.1, 0.2]}), + JointState({"position": [0.3, 0.4]}), + ) + + assert result.status == PlanningStatus.SUCCESS + assert world.projected_states + assert all(state.name == ["joint_a", "joint_b", "gripper"] for state in world.projected_states) + assert all(state.position[2] == 0.77 for state in world.projected_states) diff --git a/dimos/manipulation/planning/spec/enums.py b/dimos/manipulation/planning/spec/enums.py index 66a17ee199..3d5b668e4b 100644 --- a/dimos/manipulation/planning/spec/enums.py +++ b/dimos/manipulation/planning/spec/enums.py @@ -35,6 +35,7 @@ class IKStatus(Enum): JOINT_LIMITS = auto() COLLISION = auto() TIMEOUT = auto() + UNSUPPORTED = auto() class PlanningStatus(Enum): @@ -47,3 +48,4 @@ class PlanningStatus(Enum): INVALID_GOAL = auto() COLLISION_AT_START = auto() COLLISION_AT_GOAL = auto() + UNSUPPORTED = auto() diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index b5fea13a36..9ea84ab646 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -20,6 +20,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -28,6 +29,7 @@ import numpy as np from numpy.typing import NDArray + from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( IKResult, @@ -238,6 +240,20 @@ def solve( """Solve IK with optional collision checking.""" ... + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve planning-group-scoped pose targets.""" + ... + @runtime_checkable class PlannerSpec(Protocol): @@ -263,6 +279,18 @@ def plan_joint_path( """Plan a collision-free joint-space path.""" ... + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan a collision-free path for an ordered planning-group selection.""" + ... + def get_name(self) -> str: """Get planner name.""" ... diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index b5ac50d828..d244bfe246 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -45,7 +45,7 @@ make_global_joint_names, make_planning_group_id, ) -from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus @@ -415,6 +415,21 @@ def plan_joint_path( message="RoboPlan path found", ) + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """PlanningGroupSelection planning is not implemented by RoboPlanWorld.""" + return PlanningResult( + status=PlanningStatus.UNSUPPORTED, + message="RoboPlanWorld does not support selected planning-group joint paths", + ) + def get_name(self) -> str: """Get planner name.""" return "RoboPlan" diff --git a/openspec/changes/group-aware-ik-rrt/.openspec.yaml b/openspec/changes/group-aware-ik-rrt/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/group-aware-ik-rrt/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/group-aware-ik-rrt/README.md b/openspec/changes/group-aware-ik-rrt/README.md deleted file mode 100644 index 872d917e36..0000000000 --- a/openspec/changes/group-aware-ik-rrt/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# group-aware-ik-rrt - -Extract PR 3 from reference branch cc/spec/movegroup: make IK and RRT consume explicit planning groups and group-local targets. diff --git a/openspec/changes/group-aware-ik-rrt/design.md b/openspec/changes/group-aware-ik-rrt/design.md deleted file mode 100644 index 6187e24bb9..0000000000 --- a/openspec/changes/group-aware-ik-rrt/design.md +++ /dev/null @@ -1,27 +0,0 @@ -## Context - -IK and RRT previously operated mostly through robot-scoped assumptions. With groups, each target must map to an explicit group, local joint order, and target frame. - -## Goals / Non-Goals - -**Goals:** -- Make IK target frames come from the requested planning group. -- Make RRT accept and return group-local joint targets and paths. -- Keep collision checks and global robot-state projection correct. -- Fail clearly when a requested group lacks required pose target metadata. - -**Non-Goals:** -- Do not change `ManipulationModule` public APIs in this PR. -- Do not include visualization changes. -- Do not change control tasks. - -## Decisions - -- Prefer explicit group IDs for all new solver/planner paths. -- Robot-scoped compatibility may resolve through a unique pose-targetable group, but ambiguous robots must fail clearly. -- PinkIK and Viser-style base-pose transforms must validate that robot-scoped `base_link` is compatible with model-root assumptions. - -## Risks / Trade-offs - -- Algorithm tests can become broad quickly. Keep this PR focused on solver/planner contracts and leave module-level behavior to PR 4. -- Some optional dependencies may be unavailable locally; fake dependency tests should stay hermetic. diff --git a/openspec/changes/group-aware-ik-rrt/proposal.md b/openspec/changes/group-aware-ik-rrt/proposal.md deleted file mode 100644 index 0ea8cb9a14..0000000000 --- a/openspec/changes/group-aware-ik-rrt/proposal.md +++ /dev/null @@ -1,24 +0,0 @@ -## Why - -After world backends understand planning groups, the solver and planner layer must consume group IDs and group-local targets. This PR keeps algorithm review separate from both backend plumbing and public module/UI migration. - -## What Changes - -- Update PinkIK, Jacobian IK, and Drake optimization IK to resolve target frames from planning groups. -- Update RRT planning to handle group-local joint targets, group candidate selection, collision checking, and result projection. -- Add focused IK/planner tests. - -## Capabilities - -### New Capabilities -- `group-aware-ik-rrt`: IK and RRT planning support explicit planning groups and group-local targets. - -### Modified Capabilities -- `manipulation-planning-algorithms`: Planners and IK solvers no longer rely on robot-scoped end-effector metadata. - -## Impact - -- Base branch: PR 2 `planning-world-monitor-groups`. -- Reference implementation: `cc/spec/movegroup`. -- Primary files: `dimos/manipulation/planning/kinematics/*`, `dimos/manipulation/planning/planners/rrt_planner.py`, related tests. -- Out of scope: public `ManipulationModule` API migration and Viser UI. diff --git a/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md b/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md deleted file mode 100644 index 8d477a3035..0000000000 --- a/openspec/changes/group-aware-ik-rrt/specs/group-aware-ik-rrt/spec.md +++ /dev/null @@ -1,22 +0,0 @@ -## ADDED Requirements - -### Requirement: IK solvers must use planning-group target frames -IK solvers MUST resolve pose target frames from the requested planning group's `tip_link` rather than from robot-scoped end-effector metadata. - -#### Scenario: Pose IK targets a manipulator group -- **WHEN** a pose target is submitted for a group with `tip_link="tcp"` -- **THEN** the IK solver constrains the `tcp` frame - -### Requirement: Planners must preserve group-local joint ordering -Planners MUST accept and return joint targets in the requested group's local joint order while projecting through full robot state for collision checks. - -#### Scenario: Group joint target uses subset order -- **WHEN** a group target names a subset of robot joints -- **THEN** planning uses the correct full robot state and returns a path scoped to the group - -### Requirement: Algorithms must fail clearly for non-pose-targetable groups -Algorithms that require a pose target frame MUST return an explicit failure when the requested group has no `tip_link`. - -#### Scenario: Pose IK targets a joint-only group -- **WHEN** a pose target is submitted for a group without `tip_link` -- **THEN** the solver returns a no-solution or unsupported result with an explanatory message diff --git a/openspec/changes/group-aware-ik-rrt/tasks.md b/openspec/changes/group-aware-ik-rrt/tasks.md deleted file mode 100644 index 3a13b0a9f7..0000000000 --- a/openspec/changes/group-aware-ik-rrt/tasks.md +++ /dev/null @@ -1,19 +0,0 @@ -## 1. Extract algorithm changes - -- [ ] 1.1 Start from PR 2 and use `cc/spec/movegroup` as reference. -- [ ] 1.2 Extract PinkIK group target-frame and base-link validation changes. -- [ ] 1.3 Extract Jacobian IK group selection changes. -- [ ] 1.4 Extract Drake optimization IK group target-frame changes. -- [ ] 1.5 Extract RRT planner group-local target and projection changes. - -## 2. Tests - -- [ ] 2.1 Bring over PinkIK tests relevant to group target frames and base-link validation. -- [ ] 2.2 Bring over Jacobian IK selection tests. -- [ ] 2.3 Bring over RRT planner selection/group tests. - -## 3. Validation - -- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/kinematics/test_pink_ik.py dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py dimos/manipulation/planning/planners/test_rrt_planner_selection.py -q`. -- [ ] 3.2 Run targeted mypy on changed algorithm files. -- [ ] 3.3 Optional manual smoke: solve IK and plan a small reachable group pose; print status and path length. diff --git a/openspec/specs/group-aware-ik-rrt/spec.md b/openspec/specs/group-aware-ik-rrt/spec.md new file mode 100644 index 0000000000..e6ea4839f5 --- /dev/null +++ b/openspec/specs/group-aware-ik-rrt/spec.md @@ -0,0 +1,38 @@ +## Purpose + +Define the planning-group-aware IK and selected-joint planning algorithm contracts. + +## Requirements + +### Requirement: IK solvers must use planning-group target frames +IK solvers MUST resolve pose target frames from the requested planning group's `tip_link` rather than from robot-scoped end-effector metadata. + +`solve_pose_targets` is a pose-target set IK API: callers MAY provide one or more pose targets keyed by planning group. Backends MUST either solve the requested target set or return `IKStatus.UNSUPPORTED` when the requested target-set shape is outside that backend's capability. + +#### Scenario: Pose IK targets a manipulator group +- **WHEN** a pose target is submitted for a group with `tip_link="tcp"` +- **THEN** the IK solver constrains the `tcp` frame + +#### Scenario: Pink IK receives multiple pose targets +- **WHEN** Pink IK receives multiple pose targets for pose-targetable planning groups +- **THEN** it solves the target frames as a target set and returns selected joints in planning-group selection order + +#### Scenario: Single-target IK backend receives multiple pose targets +- **WHEN** Drake optimization IK or Jacobian IK receives more than one pose target or auxiliary planning groups +- **THEN** it returns `IKStatus.UNSUPPORTED` with an explanatory message + +### Requirement: Planners must preserve group-local joint ordering +Planners MUST accept and return joint targets in the requested group's local joint order while projecting through full robot state for collision checks. + +Selected joint planning MUST be part of the planner protocol. Planner backends that do not support selected/group planning MUST return `PlanningStatus.UNSUPPORTED` with an explanatory message rather than treating unsupported capability as a failed search. + +#### Scenario: Group joint target uses subset order +- **WHEN** a group target names a subset of robot joints +- **THEN** planning uses the correct full robot state and returns a path scoped to the group + +### Requirement: Algorithms must fail clearly for non-pose-targetable groups +Algorithms that require a pose target frame MUST return an explicit unsupported failure when the requested group has no `tip_link`. + +#### Scenario: Pose IK targets a joint-only group +- **WHEN** a pose target is submitted for a group without `tip_link` +- **THEN** the solver returns `IKStatus.UNSUPPORTED` with an explanatory message From 05c25787ab002a31f28785ca7d00d274ddd565df Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Wed, 8 Jul 2026 22:28:57 -0700 Subject: [PATCH 18/29] Plannning Group 4: Expose manipulation planning group APIs (#2801) --- .../test_manipulation_planning_groups.py | 216 ++++++ dimos/manipulation/manipulation_module.py | 710 +++++++++++++----- .../planning/examples/manipulation_client.py | 64 +- .../manipulation/planning/groups/registry.py | 13 +- .../planning/groups/test_planning_groups.py | 70 ++ dimos/manipulation/planning/groups/utils.py | 56 ++ dimos/manipulation/planning/spec/models.py | 17 + .../planning/world/drake_world.py | 44 +- .../world/test_drake_world_planning_groups.py | 65 ++ .../manipulation/test_manipulation_module.py | 45 +- dimos/manipulation/test_manipulation_unit.py | 535 +++++++++++-- .../visualization/viser/adapter.py | 44 +- dimos/manipulation/visualization/viser/gui.py | 2 - .../manipulation/visualization/viser/state.py | 1 - .../viser/test_viser_visualization.py | 47 +- dimos/robot/all_blueprints.py | 1 + dimos/robot/manipulators/test_blueprints.py | 27 +- .../manipulators/xarm/blueprints/basic.py | 21 + dimos/robot/test_all_blueprints.py | 1 + .../.openspec.yaml | 2 - .../manipulation-module-group-api/README.md | 3 - .../manipulation-module-group-api/design.md | 28 - .../manipulation-module-group-api/proposal.md | 25 - .../manipulation-module-group-api/tasks.md | 18 - .../manipulation-module-group-api/spec.md | 6 +- 25 files changed, 1663 insertions(+), 398 deletions(-) create mode 100644 dimos/e2e_tests/test_manipulation_planning_groups.py delete mode 100644 openspec/changes/manipulation-module-group-api/.openspec.yaml delete mode 100644 openspec/changes/manipulation-module-group-api/README.md delete mode 100644 openspec/changes/manipulation-module-group-api/design.md delete mode 100644 openspec/changes/manipulation-module-group-api/proposal.md delete mode 100644 openspec/changes/manipulation-module-group-api/tasks.md rename openspec/{changes/manipulation-module-group-api => }/specs/manipulation-module-group-api/spec.md (88%) diff --git a/dimos/e2e_tests/test_manipulation_planning_groups.py b/dimos/e2e_tests/test_manipulation_planning_groups.py new file mode 100644 index 0000000000..67cf767177 --- /dev/null +++ b/dimos/e2e_tests/test_manipulation_planning_groups.py @@ -0,0 +1,216 @@ +# 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. + +"""Large E2E tests for manipulation planning groups with a coordinator. + +These tests launch a real ManipulationModule + ControlCoordinator blueprint and +exercise the public planning RPCs over LCM, matching the self-hosted large-test +style used by the navigation stack. +""" + +from __future__ import annotations + +from collections.abc import Callable +import time +from typing import Any + +import pytest + +from dimos.control.coordinator import ControlCoordinator +from dimos.core.rpc_client import RPCClient +from dimos.e2e_tests.dimos_cli_call import DimosCliCall +from dimos.e2e_tests.lcm_spy import LcmSpy +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + +pytestmark = [pytest.mark.self_hosted_large] + +JOINT_STATE_TOPIC = "/coordinator_joint_state#sensor_msgs.JointState" +BLUEPRINT = "openarm-mock-planner-coordinator" + + +def _wait_for_robot_info( + client: RPCClient, + robot_name: str, + *, + timeout: float = 120.0, +) -> dict[str, Any]: + deadline = time.time() + timeout + last_error: BaseException | None = None + while time.time() < deadline: + try: + info = client.get_robot_info(robot_name) + if info and info.get("planning_groups"): + return info + except BaseException as exc: + last_error = exc + time.sleep(0.5) + raise TimeoutError(f"Timed out waiting for {robot_name!r} robot info") from last_error + + +def _wait_for_trajectory_completion( + client: RPCClient, + robot_name: str, + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + last_status: dict[str, Any] | None = None + while time.time() < deadline: + last_status = client.get_trajectory_status(robot_name) + if last_status is not None and last_status.get("state") == TrajectoryState.COMPLETED: + return + time.sleep(0.1) + raise TimeoutError(f"{robot_name!r} trajectory did not complete; last={last_status}") + + +def _wait_for_manipulation_state( + client: RPCClient, + state_name: str, + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + last_state: str | None = None + while time.time() < deadline: + last_state = client.get_state() + if last_state == state_name: + return + time.sleep(0.1) + raise TimeoutError(f"ManipulationModule did not reach {state_name}; last={last_state}") + + +def _wait_for_current_joints( + client: RPCClient, + robot_names: tuple[str, ...], + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + missing = robot_names + while time.time() < deadline: + missing = tuple( + robot_name + for robot_name in robot_names + if client.get_current_joints(robot_name) is None + ) + if not missing: + return + time.sleep(0.1) + raise TimeoutError(f"Timed out waiting for current joints from {missing}") + + +def _prepare_for_planning(client: RPCClient, robot_names: tuple[str, ...]) -> None: + client.reset() + _wait_for_manipulation_state(client, "IDLE") + _wait_for_current_joints(client, robot_names) + # Robot info and joint-state topics can become available just before the + # manipulation module finishes finalizing world monitors. Require a stable + # ready state after joint state is flowing to avoid command-readiness flakes. + time.sleep(0.25) + _wait_for_manipulation_state(client, "IDLE") + + +def _planning_group_id(info: dict[str, Any]) -> str: + groups = info["planning_groups"] + assert len(groups) == 1 + group = groups[0] + if isinstance(group, PlanningGroup): + return group.id + group_id = group["id"] + assert isinstance(group_id, str) + return group_id + + +def _offset_target(client: RPCClient, robot_name: str, delta: float) -> JointState: + current = client.get_current_joints(robot_name) + assert current is not None + return JointState(position=[position + delta for position in current]) + + +def _start_openarm_mock_planner( + start_blueprint: Callable[..., DimosCliCall], lcm_spy: LcmSpy +) -> None: + lcm_spy.save_topic(JOINT_STATE_TOPIC) + start_blueprint(BLUEPRINT) + lcm_spy.wait_for_saved_topic(JOINT_STATE_TOPIC, timeout=120.0) + + +def test_single_arm_plans_and_executes_through_control_coordinator( + lcm_spy: LcmSpy, + start_blueprint: Callable[..., DimosCliCall], +) -> None: + """Plan with one arm and execute through its trajectory task.""" + _start_openarm_mock_planner(start_blueprint, lcm_spy) + + client = RPCClient(None, ManipulationModule) + coordinator_client = RPCClient(None, ControlCoordinator) + try: + left_info = _wait_for_robot_info(client, "left_arm") + left_id = _planning_group_id(left_info) + + tasks = coordinator_client.list_tasks() + assert left_info["coordinator_task_name"] in tasks + + _prepare_for_planning(client, ("left_arm",)) + + planned = client.plan_to_joint_targets({left_id: _offset_target(client, "left_arm", 0.02)}) + assert planned, client.get_error() + assert client.has_planned_path() + assert client.execute_plan() + + _wait_for_trajectory_completion(client, "left_arm") + finally: + coordinator_client.stop_rpc_client() + client.stop_rpc_client() + + +def test_dual_arm_plans_and_dispatches_both_arms_through_control_coordinator( + lcm_spy: LcmSpy, + start_blueprint: Callable[..., DimosCliCall], +) -> None: + """Plan one generated plan over both arms and dispatch both JTC tasks.""" + _start_openarm_mock_planner(start_blueprint, lcm_spy) + + client = RPCClient(None, ManipulationModule) + coordinator_client = RPCClient(None, ControlCoordinator) + try: + left_info = _wait_for_robot_info(client, "left_arm") + right_info = _wait_for_robot_info(client, "right_arm") + left_id = _planning_group_id(left_info) + right_id = _planning_group_id(right_info) + + tasks = coordinator_client.list_tasks() + assert left_info["coordinator_task_name"] in tasks + assert right_info["coordinator_task_name"] in tasks + + _prepare_for_planning(client, ("left_arm", "right_arm")) + + planned = client.plan_to_joint_targets( + { + left_id: _offset_target(client, "left_arm", 0.02), + right_id: _offset_target(client, "right_arm", -0.02), + } + ) + assert planned, client.get_error() + assert client.has_planned_path() + assert client.execute_plan() + + _wait_for_trajectory_completion(client, "left_arm") + _wait_for_trajectory_completion(client, "right_arm") + finally: + coordinator_client.stop_rpc_client() + client.stop_rpc_client() diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 9b22dc0468..a0fed03f7b 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -24,6 +24,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from enum import Enum import threading import time @@ -45,6 +46,13 @@ create_planning_specs, create_world, ) +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import ( + filter_joint_state_to_selected_joints, + joint_target_to_global_names, + planning_group_id_from_selector, + project_global_joint_path_to_robot, +) from dimos.manipulation.planning.kinematics.config import ( ManipulationKinematicsConfig, PinkKinematicsConfig, @@ -53,9 +61,11 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus, ObstacleType from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, IKResult, JointPath, Obstacle, + PlanningGroupID, RobotName, WorldRobotID, ) @@ -71,6 +81,7 @@ from dimos.manipulation.visualization.factory import create_manipulation_visualization from dimos.manipulation.visualization.types import TargetEvaluation 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.sensor_msgs.JointState import JointState @@ -89,11 +100,10 @@ RobotRegistry: TypeAlias = dict[RobotName, RobotEntry] """Maps robot_name -> RobotEntry""" -PlannedPaths: TypeAlias = dict[RobotName, JointPath] -"""Maps robot_name -> planned joint path""" - -PlannedTrajectories: TypeAlias = dict[RobotName, JointTrajectory] -"""Maps robot_name -> planned trajectory""" +RobotInfoValue: TypeAlias = ( + str | bool | float | list[str] | list[float] | list[PlanningGroup] | None +) +RobotInfoPayload: TypeAlias = dict[str, RobotInfoValue] class ManipulationState(Enum): @@ -156,9 +166,9 @@ def __init__(self, **kwargs: Any) -> None: # Robot registry: maps robot_name -> (world_robot_id, config, trajectory_gen) self._robots: RobotRegistry = {} - # Stored path for plan/preview/execute workflow (per robot) - self._planned_paths: PlannedPaths = {} - self._planned_trajectories: PlannedTrajectories = {} + # Canonical generated plan for plan/preview/execute workflow. + # Robot-local paths and trajectories are derived from this plan on demand. + self._last_plan: GeneratedPlan | None = None # Coordinator integration (lazy initialized) self._coordinator_client: RPCClient | None = None @@ -447,7 +457,11 @@ def get_ee_pose(self, robot_name: RobotName | None = None) -> Pose | None: robot_name: Robot to query (required if multiple robots configured) """ if (robot := self._get_robot(robot_name)) and self._world_monitor: - return self._world_monitor.get_ee_pose(robot[1], joint_state=None) + try: + return self._world_monitor.get_ee_pose(robot[1], joint_state=None) + except ValueError as exc: + logger.warning("End-effector pose unavailable: %s", exc) + return None return None @rpc @@ -485,6 +499,140 @@ def _begin_planning( self._state = ManipulationState.PLANNING return robot[0], robot[1] + def _begin_group_planning(self) -> int | None: + """Check state and begin planning for explicit planning-group APIs.""" + if self._world_monitor is None: + logger.error("Planning not initialized") + return None + with self._lock: + if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): + logger.warning(f"Cannot plan: state is {self._state.name}") + return None + self._planning_epoch += 1 + self._last_plan = None + self._state = ManipulationState.PLANNING + return self._planning_epoch + + def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID: + """Return the unique pose-targetable group or raise if it is ambiguous.""" + if self._world_monitor is None: + raise ValueError("Planning not initialized") + group_id = self._world_monitor.planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError( + f"Robot '{robot_name}' has no pose-targetable planning group; " + "use an explicit planning group ID" + ) + return group_id + + def _project_plan_to_robot_paths( + self, plan: GeneratedPlan + ) -> dict[RobotName, JointPath] | None: + """Project a selected-global-joint plan into full local paths per affected robot. + + TODO: need proper trajectory parametrization + """ + if self._world_monitor is None: + return None + paths: dict[RobotName, JointPath] = {} + try: + affected = self._world_monitor.planning_groups.select(plan.group_ids).robot_names + except Exception as exc: + logger.error("Failed to resolve generated plan: %s", exc) + return None + + for robot_name in affected: + robot = self._get_robot(robot_name) + if robot is None: + return None + resolved_name, robot_id, config, _ = robot + current = self._world_monitor.get_current_joint_state(robot_id) + try: + paths[resolved_name] = project_global_joint_path_to_robot( + plan.path, + robot_name=resolved_name, + local_joint_names=config.joint_names, + current_joint_state=current, + ) + except ValueError as exc: + logger.error("Cannot project plan for '%s': %s", resolved_name, exc) + return None + return paths + + def _trajectory_from_robot_path( + self, robot_name: RobotName, path: JointPath + ) -> JointTrajectory | None: + """Generate a local trajectory for one projected robot path.""" + if len(path) < 2: + logger.error("Plan projection for '%s' has fewer than two waypoints", robot_name) + return None + robot = self._get_robot(robot_name) + if robot is None: + return None + _, _, config, traj_gen = robot + local_trajectory = traj_gen.generate([list(state.position) for state in path]) + return JointTrajectory( + joint_names=list(config.joint_names), + points=local_trajectory.points, + timestamp=local_trajectory.timestamp, + ) + + def _trajectories_for_plan( + self, plan: GeneratedPlan + ) -> dict[RobotName, JointTrajectory] | None: + """Derive robot-local trajectories from a generated plan.""" + projected = self._project_plan_to_robot_paths(plan) + if projected is None: + return None + trajectories: dict[RobotName, JointTrajectory] = {} + for robot_name, path in projected.items(): + trajectory = self._trajectory_from_robot_path(robot_name, path) + if trajectory is None: + return None + trajectories[robot_name] = trajectory + logger.info("Trajectory for %s: %.3fs", robot_name, trajectory.duration) + return trajectories + + def _plan_selected_path( + self, + group_ids: tuple[PlanningGroupID, ...], + start: JointState, + goal: JointState, + planning_epoch: int, + ) -> bool: + """Plan over an explicit planning group selection and store the result.""" + assert self._world_monitor and self._planner + result = self._planner.plan_selected_joint_path( + world=self._world_monitor.world, + selection=self._world_monitor.planning_groups.select(group_ids), + start=start, + goal=goal, + timeout=self.config.planning_timeout, + ) + if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: + logger.info("Discarding cancelled planning result") + return False + if not result.is_success(): + return self._fail(f"Planning failed: {result.status.name}") + + logger.info("Path: %d waypoints, groups=%s", len(result.path), group_ids) + plan = GeneratedPlan( + group_ids=group_ids, + path=result.path, + status=result.status, + planning_time=result.planning_time, + path_length=result.path_length, + iterations=result.iterations, + message=result.message, + ) + projected = self._project_plan_to_robot_paths(plan) + if projected is None or any(len(path) < 2 for path in projected.values()): + self._last_plan = None + return self._fail("Failed to project generated plan") + self._last_plan = plan + self._state = ManipulationState.COMPLETED + return True + def _fail(self, msg: str) -> bool: """Set FAULT state with error message.""" logger.warning(msg) @@ -509,9 +657,6 @@ def _solve_ik_for_pose( """Run the configured kinematics backend for a world-frame pose.""" assert self._world_monitor and self._kinematics - # Convert Pose to PoseStamped for the IK solver - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - target_pose = PoseStamped( frame_id="world", position=pose.position, @@ -527,27 +672,88 @@ def _solve_ik_for_pose( ) @rpc - def solve_ik( + def inverse_kinematics( + self, + pose_targets: Mapping[PlanningGroupID, PoseStamped], + auxiliary_group_ids: Sequence[PlanningGroupID] = (), + seed: JointState | None = None, + check_collision: bool = True, + ) -> IKResult: + """Solve planning-group pose targets without planning a joint path.""" + if self._kinematics is None or self._world_monitor is None: + return IKResult(status=IKStatus.NO_SOLUTION, message="Planning not initialized") + if not pose_targets: + return IKResult( + status=IKStatus.NO_SOLUTION, message="At least one pose target is required" + ) + + try: + stamped_targets = dict(pose_targets) + auxiliary_ids = tuple(auxiliary_group_ids) + group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) + target_groups = { + self._world_monitor.planning_groups.get(group_id): pose + for group_id, pose in stamped_targets.items() + } + auxiliary_groups = tuple( + self._world_monitor.planning_groups.get(group_id) for group_id in auxiliary_ids + ) + seed_state = seed + if seed_state is None: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + if not current.name and not current.position: + return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") + seed_state = filter_joint_state_to_selected_joints(current, selection.joint_names) + except (KeyError, ValueError) as exc: + return IKResult(status=IKStatus.NO_SOLUTION, message=str(exc)) + if seed_state is None: + return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") + return self._kinematics.solve_pose_targets( + world=self._world_monitor.world, + pose_targets=target_groups, + auxiliary_groups=auxiliary_groups, + seed=seed_state, + check_collision=check_collision, + ) + + @rpc + def inverse_kinematics_single( self, pose: Pose, robot_name: RobotName | None = None, - check_collision: bool = True, seed: JointState | None = None, + check_collision: bool = True, ) -> IKResult: - """Solve IK for a pose without planning a joint path. - - Args: - pose: Target end-effector pose - robot_name: Robot to solve for (required if multiple robots configured) - check_collision: Whether to reject IK candidates in collision - seed: Optional joint state to initialize local IK. Uses current state when omitted. - """ - if self._kinematics is None or self._world_monitor is None: + """Solve IK for one robot's unique pose-targetable planning group.""" + if self._world_monitor is None: return IKResult(status=IKStatus.NO_SOLUTION, message="Planning not initialized") robot = self._get_robot(robot_name) if robot is None: return IKResult(status=IKStatus.NO_SOLUTION, message="Robot not found") + selected_robot_name, _, _, _ = robot + try: + group_id = self._require_unique_pose_group_id_for_robot(selected_robot_name) + except ValueError as exc: + return IKResult(status=IKStatus.NO_SOLUTION, message=str(exc)) + target_pose = PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + return self.inverse_kinematics( + {group_id: target_pose}, seed=seed, check_collision=check_collision + ) + @rpc + def solve_ik( + self, + pose: Pose, + robot_name: RobotName | None = None, + check_collision: bool = True, + seed: JointState | None = None, + ) -> IKResult: + """Compatibility wrapper for inverse_kinematics_single().""" with self._lock: if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): return IKResult( @@ -555,14 +761,12 @@ def solve_ik( message=f"Cannot solve IK while state is {self._state.name}", ) self._state = ManipulationState.PLANNING - - _, robot_id, _, _ = robot - seed_state = seed or self._world_monitor.get_current_joint_state(robot_id) - if seed_state is None: - self._state = ManipulationState.IDLE - return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") - - result = self._solve_ik_for_pose(robot_id, pose, seed_state, check_collision) + result = self.inverse_kinematics_single( + pose, + robot_name=robot_name, + seed=seed, + check_collision=check_collision, + ) self._state = ManipulationState.COMPLETED if result.is_success() else ManipulationState.IDLE if result.is_success(): logger.info(f"IK solved, error: {result.position_error:.4f}m") @@ -570,127 +774,151 @@ def solve_ik( @rpc def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: - """Plan motion to pose. Use preview_path() then execute(). + """Plan motion to pose. Use preview_plan() then execute(). Args: pose: Target end-effector pose robot_name: Robot to plan for (required if multiple robots configured) """ - if self._kinematics is None or (r := self._begin_planning(robot_name)) is None: + if self._kinematics is None or self._world_monitor is None: + return False + robot = self._get_robot(robot_name) + if robot is None: + return False + selected_robot_name, _, _, _ = robot + try: + group_id = self._require_unique_pose_group_id_for_robot(selected_robot_name) + except ValueError as exc: + logger.warning("Pose planning unavailable: %s", exc) return False - robot_name, robot_id = r - planning_epoch = self._planning_epoch - assert self._world_monitor # guaranteed by _begin_planning + return self.plan_to_pose_targets({group_id: pose}) - current = self._world_monitor.get_current_joint_state(robot_id) - if current is None: - return self._fail("No joint state") + @rpc + def plan_to_pose_targets( + self, + pose_targets: Mapping[PlanningGroupID | PlanningGroup, Pose], + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> bool: + """Plan to one or more group pose targets with optional auxiliary groups.""" + if self._world_monitor is None or self._kinematics is None: + return False + if not pose_targets: + return self._fail("At least one pose target is required") + + stamped_targets = { + planning_group_id_from_selector(group): PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + for group, pose in pose_targets.items() + } + auxiliary_ids = tuple(planning_group_id_from_selector(group) for group in auxiliary_groups) + group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) + planning_epoch = self._begin_group_planning() + if planning_epoch is None: + return False - ik = self._solve_ik_for_pose(robot_id, pose, current, check_collision=True) + try: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + except Exception as exc: + return self._fail(f"Failed to resolve planning groups: {exc}") + + ik = self.inverse_kinematics( + pose_targets=stamped_targets, + auxiliary_group_ids=auxiliary_ids, + seed=start, + ) if not ik.is_success() or ik.joint_state is None: return self._fail(f"IK failed: {ik.status.name}") - logger.info(f"IK solved, error: {ik.position_error:.4f}m") - return self._plan_path_only(robot_name, robot_id, ik.joint_state, planning_epoch) + return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) @rpc def plan_to_joints(self, joints: JointState, robot_name: RobotName | None = None) -> bool: - """Plan motion to joint config. Use preview_path() then execute(). + """Plan motion to joint config. Use preview_plan() then execute(). Args: joints: Target joint state (names + positions) robot_name: Robot to plan for (required if multiple robots configured) """ - if (r := self._begin_planning(robot_name)) is None: + robot = self._get_robot(robot_name) + if robot is None: return False - robot_name, robot_id = r - planning_epoch = self._planning_epoch - logger.info(f"Planning to joints for {robot_name}: {[f'{j:.3f}' for j in joints.position]}") - return self._plan_path_only(robot_name, robot_id, joints, planning_epoch) - - def _plan_path_only( - self, - robot_name: RobotName, - robot_id: WorldRobotID, - goal: JointState, - planning_epoch: int, - ) -> bool: - """Plan path from current position to goal, store result.""" - assert self._world_monitor and self._planner # guaranteed by _begin_planning - self._dismiss_preview(robot_id) - start = self._world_monitor.get_current_joint_state(robot_id) - if start is None: - return self._fail("No joint state") - - # Trim goal to planner DOF (e.g. strip gripper joint from coordinator state) - planner_dof = len(start.position) - if len(goal.position) > planner_dof: - goal = JointState( - name=list(goal.name[:planner_dof]) if goal.name else [], - position=list(goal.position[:planner_dof]), - ) - - result = self._planner.plan_joint_path( - world=self._world_monitor.world, - robot_id=robot_id, - start=start, - goal=goal, - timeout=self.config.planning_timeout, + selected_robot_name, _, _, _ = robot + logger.info( + f"Planning to joints for {selected_robot_name}: {[f'{j:.3f}' for j in joints.position]}" ) - if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: - logger.info("Discarding cancelled planning result") + if self._world_monitor is None: return False - if not result.is_success(): - return self._fail(f"Planning failed: {result.status.name}") - - logger.info(f"Path: {len(result.path)} waypoints") - self._planned_paths[robot_name] = result.path - - _, _, traj_gen = self._robots[robot_name] - # Convert JointState path to list of position lists for trajectory generator - traj = traj_gen.generate([list(state.position) for state in result.path]) - self._planned_trajectories[robot_name] = traj - logger.info(f"Trajectory: {traj.duration:.3f}s") - - self._state = ManipulationState.COMPLETED - return True + group_id = self._world_monitor.planning_groups.default_group_id_for_robot( + selected_robot_name + ) + if group_id is None: + logger.error( + "Robot '%s' has no unique default planning group; use explicit group APIs", + selected_robot_name, + ) + return False + return self.plan_to_joint_targets({group_id: joints}) @rpc - def preview_path( - self, - duration: float | None = None, - robot_name: RobotName | None = None, - target_fps: float = 30.0, + def plan_to_joint_targets( + self, joint_targets: Mapping[PlanningGroupID | PlanningGroup, JointState] ) -> bool: - """Preview the planned path in the visualizer. - - Args: - duration: Total animation duration in seconds. Uses trajectory duration if None. - robot_name: Robot to preview (required if multiple robots configured) - target_fps: Nominal preview update rate. Set <= 0 to use planned waypoints directly. - """ - if self._world_monitor is None: + """Plan to joint targets keyed by planning group.""" + if self._world_monitor is None or self._planner is None: return False + if not joint_targets: + return self._fail("At least one joint target is required") - robot = self._get_robot(robot_name) - if robot is None: + group_ids = tuple( + dict.fromkeys(planning_group_id_from_selector(group) for group in joint_targets) + ) + planning_epoch = self._begin_group_planning() + if planning_epoch is None: return False - robot_name, robot_id, _, _ = robot - planned_path = self._planned_paths.get(robot_name) - if planned_path is None or len(planned_path) == 0: - logger.warning(f"No planned path to preview for {robot_name}") + try: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + except Exception as exc: + return self._fail(f"Failed to resolve planning groups: {exc}") + + goal_names: list[str] = [] + goal_positions: list[float] = [] + for group, target in joint_targets.items(): + group_id = planning_group_id_from_selector(group) + try: + target_group = self._world_monitor.planning_groups.get(group_id) + target_global = joint_target_to_global_names(target_group, target) + except (KeyError, ValueError) as exc: + logger.error(str(exc)) + return self._fail(f"Invalid joint target for '{group_id}'") + goal_names.extend(target_global.name) + goal_positions.extend(target_global.position) + + goal = JointState(name=goal_names, position=goal_positions) + return self._plan_selected_path(group_ids, start, goal, planning_epoch) + + def _preview_projected_path( + self, + robot_id: WorldRobotID, + planned_path: JointPath, + trajectory: JointTrajectory, + duration: float | None, + target_fps: float, + ) -> bool: + """Animate a projected local path with optional trajectory interpolation.""" + if self._world_monitor is None or len(planned_path) == 0: return False - if duration is None: - trajectory = self._planned_trajectories.get(robot_name) - animation_duration = trajectory.duration if trajectory is not None else 3.0 - else: - trajectory = self._planned_trajectories.get(robot_name) - animation_duration = duration - + animation_duration = duration if duration is not None else trajectory.duration interpolated = list(planned_path) - if trajectory is not None and target_fps > 0 and animation_duration > 0: + if target_fps > 0 and animation_duration > 0: times = np.array( [point.time_from_start for point in trajectory.points], dtype=np.float64 ) @@ -712,6 +940,61 @@ def preview_path( self._world_monitor.animate_path(robot_id, interpolated, animation_duration) return True + @rpc + def preview_path( + self, + duration: float | None = None, + robot_name: RobotName | None = None, + target_fps: float = 30.0, + ) -> bool: + """Compatibility wrapper for preview_plan(). + + Args: + duration: Total animation duration in seconds. Uses trajectory duration if None. + robot_name: Robot to preview (required if multiple robots configured) + target_fps: Nominal preview update rate. Set <= 0 to use planned waypoints directly. + """ + return self.preview_plan(None, duration, robot_name, target_fps) + + @rpc + def preview_plan( + self, + plan: GeneratedPlan | None = None, + duration: float | None = None, + robot_name: RobotName | None = None, + target_fps: float = 30.0, + ) -> bool: + """Preview a generated plan in the visualizer.""" + plan = plan or self._last_plan + if plan is None or not plan.path: + logger.warning("No generated plan to preview") + return False + try: + assert self._world_monitor is not None + affected = list(self._world_monitor.planning_groups.select(plan.group_ids).robot_names) + except Exception as exc: + logger.error("Generated plan cannot be resolved: %s", exc) + return False + if robot_name is not None: + if robot_name not in affected: + logger.error("Generated plan does not affect robot '%s'", robot_name) + return False + affected = [robot_name] + projected = self._project_plan_to_robot_paths(plan) + if projected is None: + return False + for name in affected: + robot = self._get_robot(name) + if robot is None or (planned_path := projected.get(name)) is None: + return False + _, robot_id, _, _ = robot + trajectory = self._trajectory_from_robot_path(name, planned_path) + if trajectory is None or not self._preview_projected_path( + robot_id, planned_path, trajectory, duration, target_fps + ): + return False + return True + @rpc def has_planned_path(self) -> bool: """Check if there's a planned path ready. @@ -719,13 +1002,7 @@ def has_planned_path(self) -> bool: Returns: True if a path is planned and ready """ - robot = self._get_robot() - if robot is None: - return False - robot_name, _, _, _ = robot - - path = self._planned_paths.get(robot_name) - return path is not None and len(path) > 0 + return self._last_plan is not None and bool(self._last_plan.path) @rpc def get_visualization_url(self) -> str | None: @@ -747,13 +1024,7 @@ def clear_planned_path(self) -> bool: """ if self._world_monitor is None: return False - robot = self._get_robot() - if robot is None: - return False - robot_name, _, _, _ = robot - - self._planned_paths.pop(robot_name, None) - self._planned_trajectories.pop(robot_name, None) + self._last_plan = None return True @rpc @@ -766,7 +1037,23 @@ def list_robots(self) -> list[str]: return list(self._robots.keys()) @rpc - def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: + def list_planning_groups(self) -> list[PlanningGroup]: + """Return all configured planning groups.""" + if self._world_monitor is None: + return [] + return list(self._world_monitor.planning_groups.list()) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + """Return the named robot's current local joint state with names.""" + if self._world_monitor is None: + return None + robot_id = self.robot_id_for_name(robot_name) + if robot_id is None: + return None + return self._world_monitor.get_current_joint_state(robot_id) + + @rpc + def get_robot_info(self, robot_name: RobotName | None = None) -> RobotInfoPayload | None: """Get information about a robot. Args: @@ -780,12 +1067,22 @@ def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, Any] return None robot_name, robot_id, config, _ = robot + planning_groups = ( + list(self._world_monitor.planning_groups.groups_for_robot(robot_name)) + if self._world_monitor is not None + else [] + ) + try: + end_effector_link = config.end_effector_link + except ValueError: + end_effector_link = None return { "name": config.name, "world_robot_id": robot_id, "joint_names": config.joint_names, - "end_effector_link": config.end_effector_link, + "planning_groups": planning_groups, + "end_effector_link": end_effector_link, "base_link": config.base_link, "max_velocity": config.max_velocity, "max_acceleration": config.max_acceleration, @@ -908,20 +1205,6 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu "collision_free": collision_free, } - def get_planned_path(self, robot_name: RobotName) -> JointPath | None: - """Return a copy of the stored planned path for visualization.""" - path = self._planned_paths.get(robot_name) - if path is None: - return None - return [JointState(point) for point in path] - - def get_planned_trajectory_duration(self, robot_name: RobotName) -> float | None: - """Return the stored planned trajectory duration for visualization.""" - trajectory = self._planned_trajectories.get(robot_name) - if trajectory is None: - return None - return float(trajectory.duration) - @rpc def set_init_joints(self, joint_state: JointState, robot_name: RobotName | None = None) -> bool: """Set the init joint state. @@ -1007,42 +1290,109 @@ def _translate_trajectory_to_coordinator( timestamp=trajectory.timestamp, ) - @rpc - def execute(self, robot_name: RobotName | None = None) -> bool: - """Execute planned trajectory via ControlCoordinator.""" + def _execute_robot_trajectory(self, robot_name: RobotName, trajectory: JointTrajectory) -> bool: + """Execute one robot's local trajectory through ControlCoordinator.""" if (robot := self._get_robot(robot_name)) is None: return False - robot_name, _, config, _ = robot + resolved_name, _, config, _ = robot - if (traj := self._planned_trajectories.get(robot_name)) is None: - logger.warning("No planned trajectory") - return False if not config.coordinator_task_name: - logger.error(f"No coordinator_task_name for '{robot_name}'") + logger.error(f"No coordinator_task_name for '{resolved_name}'") return False if (client := self._get_coordinator_client()) is None: logger.error("No coordinator client") return False - translated = self._translate_trajectory_to_coordinator(traj, config) + translated = self._translate_trajectory_to_coordinator(trajectory, config) logger.info( f"Executing: task='{config.coordinator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" ) - self._state = ManipulationState.EXECUTING result = client.task_invoke( config.coordinator_task_name, "execute", {"trajectory": translated} ) if result: logger.info("Trajectory accepted") - self._state = ManipulationState.COMPLETED return True - else: - return self._fail("Coordinator rejected trajectory") + return self._fail("Coordinator rejected trajectory") + + @rpc + def execute(self, robot_name: RobotName | None = None) -> bool: + """Compatibility wrapper for execute_plan().""" + return self.execute_plan(robot_name=robot_name) + + @rpc + def execute_plan( + self, plan: GeneratedPlan | None = None, robot_name: RobotName | None = None + ) -> bool: + """Execute a generated planning-group plan through affected trajectory tasks.""" + if self._world_monitor is None: + logger.error("Planning not initialized") + return False + plan = plan or self._last_plan + if plan is None or not plan.path: + logger.warning("No generated plan") + return False + if (client := self._get_coordinator_client()) is None: + logger.error("No coordinator client") + return False + try: + affected = list(self._world_monitor.planning_groups.select(plan.group_ids).robot_names) + except Exception as exc: + return self._fail(f"Failed to resolve generated plan: {exc}") + if robot_name is not None: + if robot_name not in affected: + logger.warning("No planned path for '%s'", robot_name) + return False + affected = [robot_name] + trajectories = self._trajectories_for_plan(plan) + if trajectories is None: + return self._fail("Failed to project generated plan") + + self._state = ManipulationState.EXECUTING + for robot_name in affected: + robot = self._get_robot(robot_name) + if robot is None: + return self._fail(f"Robot '{robot_name}' not found") + _, _, config, _ = robot + if not config.coordinator_task_name: + logger.error(f"No coordinator_task_name for '{robot_name}'") + return self._fail(f"No coordinator_task_name for '{robot_name}'") + traj = trajectories.get(robot_name) + if traj is None: + logger.warning("No planned trajectory for '%s'", robot_name) + return self._fail(f"No planned trajectory for '{robot_name}'") + translated = self._translate_trajectory_to_coordinator(traj, config) + logger.info( + "Executing: task='%s', %d pts, %.2fs", + config.coordinator_task_name, + len(translated.points), + translated.duration, + ) + result = client.task_invoke( + config.coordinator_task_name, "execute", {"trajectory": translated} + ) + if not result: + return self._fail("Coordinator rejected trajectory") + + logger.info("Trajectory accepted") + self._state = ManipulationState.COMPLETED + return True @rpc def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: """Get trajectory execution status via coordinator task_invoke.""" + last_plan = self._last_plan + if robot_name is None and last_plan is not None and last_plan.path: + if self._world_monitor is None: + return None + statuses = { + name: self.get_trajectory_status(name) + for name in self._world_monitor.planning_groups.select( + last_plan.group_ids + ).robot_names + } + return {"robots": statuses} if (robot := self._get_robot(robot_name)) is None: return None _, _, config, _ = robot @@ -1200,18 +1550,39 @@ def _wait_for_trajectory_completion( Returns: True if trajectory completed successfully """ + last_plan = self._last_plan + if robot_name is None and last_plan is not None and last_plan.path: + try: + assert self._world_monitor is not None + affected = self._world_monitor.planning_groups.select( + last_plan.group_ids + ).robot_names + except Exception as exc: + logger.warning("Failed to resolve generated plan while waiting: %s", exc) + return False + return all( + self._wait_for_trajectory_completion(name, timeout, poll_interval) + for name in affected + ) + robot = self._get_robot(robot_name) if robot is None: return True rname, _, config, _ = robot client = self._get_coordinator_client() + trajectory: JointTrajectory | None = None + if last_plan is not None and last_plan.path: + projected = self._project_plan_to_robot_paths(last_plan) + if projected is not None and (path := projected.get(rname)) is not None: + trajectory = self._trajectory_from_robot_path(rname, path) if client is None or not config.coordinator_task_name: # No coordinator — wait for trajectory duration as fallback - traj = self._planned_trajectories.get(rname) - if traj is not None: - logger.info(f"No coordinator status — waiting {traj.duration:.1f}s for trajectory") - time.sleep(traj.duration + 0.5) + if trajectory is not None: + logger.info( + f"No coordinator status — waiting {trajectory.duration:.1f}s for trajectory" + ) + time.sleep(trajectory.duration + 0.5) return True # Poll task state via task_invoke @@ -1233,9 +1604,8 @@ def _wait_for_trajectory_completion( return True except Exception: # Fallback: wait for trajectory duration - traj = self._planned_trajectories.get(rname) - if traj is not None: - remaining = traj.duration - (time.time() - start) + if trajectory is not None: + remaining = trajectory.duration - (time.time() - start) if remaining > 0: logger.info(f"Status poll failed — waiting {remaining:.1f}s for trajectory") time.sleep(remaining + 0.5) diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index 1185f28f21..6bc7582239 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -25,10 +25,14 @@ Available functions: joints() Get current joint positions ee() Get end-effector pose + groups() List explicit planning groups state() Get module state (IDLE, PLANNING, EXECUTING, ...) ik_pose(x,y,z, seed_joints=None) Solve IK only, without path planning + ik_group_pose(group_id,x,y,z) Solve IK for an explicit planning group plan(joints) Plan to joint configuration, e.g. plan([0.1]*7) + plan_group(group_id,joints) Plan to an explicit planning-group joint target plan_pose(x,y,z) Plan to Cartesian pose + plan_group_pose(group_id,x,y,z) Plan to an explicit planning-group pose target preview(duration=None) Preview planned path in Meshcat execute() Execute planned trajectory via coordinator home() Move to home position @@ -50,8 +54,10 @@ from dimos.core.rpc_client import RPCClient from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import IKResult 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.sensor_msgs.JointState import JointState @@ -80,6 +86,21 @@ def plan(target_joints: list[float], robot_name: str | None = None) -> bool: return _client.plan_to_joints(js, robot_name) +def groups() -> list[PlanningGroup]: + """List explicit planning groups available for group APIs.""" + return _client.list_planning_groups() + + +def plan_group(group_id: str, target_joints: list[float] | JointState) -> bool: + """Plan to a joint target for an explicit planning group.""" + target = ( + target_joints + if isinstance(target_joints, JointState) + else JointState(position=target_joints) + ) + return _client.plan_to_joint_targets({group_id: target}) + + def _make_target_pose( x: float, y: float, @@ -124,7 +145,6 @@ def ik_pose( pitch: float | None = None, yaw: float | None = None, robot_name: str | None = None, - check_collision: bool = True, seed_joints: list[float] | JointState | None = None, ) -> IKResult: """Solve IK for a Cartesian pose without path planning. @@ -137,13 +157,32 @@ def ik_pose( pitch: Optional target pitch. Preserves current orientation if omitted. yaw: Optional target yaw. Preserves current orientation if omitted. robot_name: Robot to solve for when multiple robots are configured. - check_collision: Whether to reject IK candidates in collision. seed_joints: Optional initial joint configuration for local IK. Pass either a list of joint positions in robot joint order or a named JointState. """ target = _make_target_pose(x, y, z, roll, pitch, yaw, robot_name) seed = _make_seed_joint_state(seed_joints, robot_name) - return _client.solve_ik(target, robot_name, check_collision, seed) + return _client.inverse_kinematics_single(target, robot_name, seed) + + +def ik_group_pose( + group_id: str, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, + seed: JointState | None = None, +) -> IKResult: + """Solve IK for an explicit planning group pose target.""" + target = _make_target_pose(x, y, z, roll, pitch, yaw) + stamped = PoseStamped( + frame_id="world", + position=target.position, + orientation=target.orientation, + ) + return _client.inverse_kinematics({group_id: stamped}, seed=seed) def plan_pose( @@ -160,13 +199,26 @@ def plan_pose( return _client.plan_to_pose(target, robot_name) +def plan_group_pose( + group_id: str, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, +) -> bool: + """Plan to a Cartesian pose for an explicit planning group.""" + target = _make_target_pose(x, y, z, roll, pitch, yaw) + return _client.plan_to_pose_targets({group_id: target}) + + def preview( duration: float | None = None, robot_name: str | None = None, - target_fps: float = 30.0, ) -> bool: - """Preview planned path in Meshcat.""" - return _client.preview_path(duration, robot_name, target_fps) + """Preview the last generated plan in the visualizer.""" + return _client.preview_plan(None, duration, robot_name) def execute(robot_name: str | None = None) -> bool: diff --git a/dimos/manipulation/planning/groups/registry.py b/dimos/manipulation/planning/groups/registry.py index c5887a090e..fb982b3562 100644 --- a/dimos/manipulation/planning/groups/registry.py +++ b/dimos/manipulation/planning/groups/registry.py @@ -89,9 +89,18 @@ def groups_for_robot(self, robot_name: RobotName) -> tuple[PlanningGroup, ...]: return tuple(self._groups_by_robot.get(robot_name, ())) def default_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: - """Return the generated fallback group ID for robot-scoped wrappers.""" + """Return the group ID used by robot-scoped joint wrappers. + + Prefer the generated whole-robot fallback group. If a robot only has one + configured planning group, use that group as the unambiguous fallback. + """ group_id = make_planning_group_id(robot_name, FALLBACK_PLANNING_GROUP_NAME) - return group_id if group_id in self._groups else None + if group_id in self._groups: + return group_id + robot_groups = self.groups_for_robot(robot_name) + if len(robot_groups) == 1: + return robot_groups[0].id + return None def primary_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: """Return the unique pose-targetable group ID for robot-scoped wrappers.""" diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 8b399354a6..5c18011a62 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -36,6 +36,7 @@ joint_target_to_global_names, matching_global_joint_name, planning_group_id_from_selector, + project_global_joint_path_to_robot, ) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -364,6 +365,75 @@ def test_registry_preserves_order_and_exposes_defaults() -> None: assert registry.default_group_id_for_robot("missing") is None +def test_registry_uses_single_group_as_robot_scoped_default() -> None: + registry = PlanningGroupRegistry( + [ + _robot_config( + "solo", + planning_groups=[PlanningGroupDefinition("arm", ("joint1",), "base", "tool")], + ), + _robot_config( + "multi", + planning_groups=[ + PlanningGroupDefinition("arm", ("joint1",), "base", "tool"), + PlanningGroupDefinition("gripper", ("joint2",), "tool"), + ], + ), + ] + ) + + assert registry.default_group_id_for_robot("solo") == "solo/arm" + assert registry.default_group_id_for_robot("multi") is None + + +def test_project_global_joint_path_to_robot_overlays_selected_joints() -> None: + path = [ + JointState(name=["robot/joint1", "robot/joint3"], position=[0.1, 0.3]), + JointState(name=["robot/joint1", "robot/joint3"], position=[0.2, 0.4]), + ] + current = JointState(name=["joint1", "joint2", "joint3"], position=[0.0, 0.5, 0.0]) + + projected = project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2", "joint3"), + current_joint_state=current, + ) + + assert [point.name for point in projected] == [ + ["joint1", "joint2", "joint3"], + ["joint1", "joint2", "joint3"], + ] + assert [point.position for point in projected] == [[0.1, 0.5, 0.3], [0.2, 0.5, 0.4]] + + +def test_project_global_joint_path_to_robot_rejects_inconsistent_path() -> None: + path = [ + JointState(name=["robot/joint1"], position=[0.1]), + JointState(name=["robot/joint2"], position=[0.2]), + ] + + with pytest.raises(ValueError, match="inconsistent waypoint joint names"): + project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2"), + current_joint_state=JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + + +def test_project_global_joint_path_to_robot_requires_current_non_selected_joints() -> None: + path = [JointState(name=["robot/joint1"], position=[0.1])] + + with pytest.raises(ValueError, match="missing joint 'joint2'"): + project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2"), + current_joint_state=JointState(name=["joint1"], position=[0.0]), + ) + + def test_registry_rejects_duplicate_robot_and_unknown_group() -> None: registry = PlanningGroupRegistry([_robot_config()]) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 283681adcc..5d9161e5bd 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -23,12 +23,15 @@ assert_global_joint_names, assert_local_joint_names, is_global_joint_name, + make_global_joint_names, ) from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import ( GlobalJointName, + JointPath, LocalModelJointName, PlanningGroupID, + RobotName, ) from dimos.msgs.sensor_msgs.JointState import JointState @@ -139,6 +142,59 @@ def joint_target_to_global_names( return JointState(name=list(group.joint_names), position=global_positions) +def project_global_joint_path_to_robot( + path: Sequence[JointState], + *, + robot_name: RobotName, + local_joint_names: Sequence[LocalModelJointName], + current_joint_state: JointState | None, +) -> JointPath: + """Project a selected-global-joint path into one robot's local joint path.""" + if not path: + return [] + + selected_joint_names = tuple(path[0].name) + assert_global_joint_names(selected_joint_names) + if any( + len(waypoint.name) != len(waypoint.position) or tuple(waypoint.name) != selected_joint_names + for waypoint in path + ): + raise ValueError("inconsistent waypoint joint names") + + selected_joint_indices = dict( + zip(selected_joint_names, range(len(selected_joint_names)), strict=True) + ) + selected_joint_set = set(selected_joint_names) + waypoint_positions = [[float(position) for position in waypoint.position] for waypoint in path] + current_by_name = ( + dict(zip(current_joint_state.name, current_joint_state.position, strict=False)) + if current_joint_state is not None + else {} + ) + global_joint_names = make_global_joint_names(robot_name, tuple(local_joint_names)) + joint_pairs = list(zip(local_joint_names, global_joint_names, strict=True)) + try: + base_positions = [ + 0.0 if global_name in selected_joint_set else float(current_by_name[local_name]) + for local_name, global_name in joint_pairs + ] + except KeyError as exc: + raise ValueError(f"missing joint '{exc.args[0]}'") from exc + + overlay_indices = [ + (local_index, selected_joint_indices[global_name]) + for local_index, (_, global_name) in enumerate(joint_pairs) + if global_name in selected_joint_indices + ] + local_path: JointPath = [] + for waypoint_positions_by_joint in waypoint_positions: + projected_positions = base_positions.copy() + for local_index, selected_index in overlay_indices: + projected_positions[local_index] = waypoint_positions_by_joint[selected_index] + local_path.append(JointState(name=list(local_joint_names), position=projected_positions)) + return local_path + + def joint_state_to_ordered_positions( joint_state: JointState, *, diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index e746bf2bc2..2dfec92e2d 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -150,6 +150,23 @@ def is_success(self) -> bool: return self.status == PlanningStatus.SUCCESS +@dataclass +class GeneratedPlan: + """Canonical selected-planning-group path exposed by ManipulationModule.""" + + group_ids: tuple[PlanningGroupID, ...] + path: list[JointState] = field(default_factory=list) + status: PlanningStatus = PlanningStatus.NO_SOLUTION + planning_time: float = 0.0 + path_length: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if the generated plan was successful.""" + return self.status == PlanningStatus.SUCCESS + + @dataclass class CollisionObjectMessage: """Message for adding/updating/removing obstacles. diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index dae0d3f73c..2b4ad6f971 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -22,6 +22,7 @@ from pathlib import Path from threading import RLock, current_thread from typing import TYPE_CHECKING, Any +import xml.etree.ElementTree as ET import numpy as np @@ -273,7 +274,7 @@ def _load_model(self, config: RobotModelConfig) -> Any: xacro_args=config.xacro_args, convert_meshes=config.auto_convert_meshes, ) - prepared_path_obj = Path(prepared_path) + prepared_path_obj = self._strip_world_base_joint(Path(prepared_path), config) # Register package paths (not applicable to MJCF) if config.package_paths: @@ -288,9 +289,48 @@ def _load_model(self, config: RobotModelConfig) -> Any: model_instances = self._parser.AddModels(prepared_path_obj) if not model_instances: - raise ValueError(f"Failed to parse model: {prepared_path}") + raise ValueError(f"Failed to parse model: {prepared_path_obj}") return model_instances[0] + @staticmethod + def _strip_world_base_joint(model_path: Path, config: RobotModelConfig) -> Path: + if model_path.suffix != ".urdf": + return model_path + + tree = ET.parse(model_path) + root = tree.getroot() + joints = root.findall("joint") + joints_to_remove = [ + joint + for joint in joints + if joint.get("type") == "fixed" + and (parent := joint.find("parent")) is not None + and (child := joint.find("child")) is not None + and parent.get("link") == "world" + and child.get("link") == config.base_link + ] + + if not joints_to_remove: + return model_path + + for joint in joints_to_remove: + root.remove(joint) + + if not any( + element is not None and element.get("link") == "world" + for joint in root.findall("joint") + for element in (joint.find("parent"), joint.find("child")) + ): + for link in list(root.findall("link")): + if link.get("name") == "world": + root.remove(link) + + stripped_path = model_path.with_name( + f"{model_path.stem}_config_base_pose{model_path.suffix}" + ) + tree.write(stripped_path, encoding="utf-8", xml_declaration=True) + return stripped_path + def _weld_base_if_needed(self, config: RobotModelConfig, model_instance: Any) -> None: """Weld robot base to world if not already welded in URDF.""" base_body = self._plant.GetBodyByName(config.base_link, model_instance) diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index c981c76896..de4c31ee95 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -53,6 +53,33 @@ def _write_urdf(path: Path) -> None: ) +def _write_urdf_with_world_base_joint(path: Path) -> None: + path.write_text( + """ + + + + + + + + + + + + + + + + + + + + +""" + ) + + def _config( path: Path, groups: list[PlanningGroupDefinition], joints: list[str] | None = None ) -> RobotModelConfig: @@ -130,6 +157,44 @@ def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Pat assert world.get_jacobian(ctx, robot_id).shape == (6, 2) +@requires_drake +def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( + tmp_path: Path, +) -> None: + urdf = tmp_path / "robot_with_world.urdf" + _write_urdf_with_world_base_joint(urdf) + world = DrakeWorld(enable_viz=False) + left_id = world.add_robot( + RobotModelConfig( + name="left_arm", + model_path=urdf, + base_pose=PoseStamped(position=[0, 0.5, 0], orientation=[0, 0, 0, 1]), + joint_names=["joint1", "joint2"], + base_link="base_link", + planning_groups=[_arm_group("joint1", "joint2")], + ) + ) + right_id = world.add_robot( + RobotModelConfig( + name="right_arm", + model_path=urdf, + base_pose=PoseStamped(position=[0, -0.5, 0], orientation=[0, 0, 0, 1]), + joint_names=["joint1", "joint2"], + base_link="base_link", + planning_groups=[_arm_group("joint1", "joint2")], + ) + ) + world.finalize() + ctx = world.get_live_context() + + left_base_pose = world.get_link_pose(ctx, left_id, "base_link") + right_base_pose = world.get_link_pose(ctx, right_id, "base_link") + + assert left_base_pose[1, 3] == pytest.approx(0.5) + assert right_base_pose[1, 3] == pytest.approx(-0.5) + assert left_base_pose[1, 3] != pytest.approx(right_base_pose[1, 3]) + + @requires_drake def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index e0958afb8c..069a2adc4c 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -169,10 +169,32 @@ def test_plan_to_joints(self, module, joint_state_zeros): assert module._state == ManipulationState.COMPLETED assert module.has_planned_path() is True - assert "test_arm" in module._planned_trajectories - traj = module._planned_trajectories["test_arm"] + assert module._last_plan is not None + projected = module._project_plan_to_robot_paths(module._last_plan) + assert projected is not None + path = projected["test_arm"] + assert len(path) > 1 + + traj = module._trajectory_from_robot_path("test_arm", path) + assert traj is not None assert len(traj.points) > 1 assert traj.duration > 0 + assert module._last_plan.group_ids == ("test_arm/manipulator",) + + def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): + """Test planning to an explicit planning-group joint target.""" + module._on_joint_state(joint_state_zeros) + + success = module.plan_to_joint_targets( + {"test_arm/manipulator": JointState(position=[0.05] * 7)} + ) + + assert success is True + assert module._state == ManipulationState.COMPLETED + assert module._last_plan is not None + assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module.has_planned_path() is True + assert module._project_plan_to_robot_paths(module._last_plan) is not None def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" @@ -200,6 +222,12 @@ def test_robot_info(self, module): assert info["end_effector_link"] == "link7" assert info["coordinator_task_name"] == "traj_arm" assert info["has_joint_name_mapping"] is True + groups = info["planning_groups"] + assert len(groups) == 1 + assert groups[0].id == "test_arm/manipulator" + + all_groups = module.list_planning_groups() + assert [group.id for group in all_groups] == ["test_arm/manipulator"] def test_ee_pose(self, module, joint_state_zeros): """Test getting end-effector pose.""" @@ -219,13 +247,17 @@ def test_trajectory_name_translation(self, module, joint_state_zeros): success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True - traj = module._planned_trajectories["test_arm"] + assert module._last_plan is not None + projected = module._project_plan_to_robot_paths(module._last_plan) + assert projected is not None + path = projected["test_arm"] + traj = module._trajectory_from_robot_path("test_arm", path) + assert traj is not None robot_config = module._robots["test_arm"][1] translated = module._translate_trajectory_to_coordinator(traj, robot_config) - for name in translated.joint_names: - assert name.startswith("arm_") # Should have arm_ prefix + assert translated.joint_names == list(robot_config.joint_name_mapping.keys()) @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @@ -260,7 +292,8 @@ def test_execute_with_mock_coordinator(self, module, joint_state_zeros): trajectory = kwargs["trajectory"] assert len(trajectory.points) > 1 # Joint names should be translated - assert all(n.startswith("arm_") for n in trajectory.joint_names) + robot_config = module._robots["test_arm"][1] + assert trajectory.joint_names == list(robot_config.joint_name_mapping.keys()) def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): """Test handling of coordinator rejection.""" diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index c4adaf6212..3e27bfef7a 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -29,11 +29,17 @@ ManipulationState, ) from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.enums import IKStatus -from dimos.manipulation.planning.spec.models import IKResult, PlanningSceneInfo +from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus +from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, + IKResult, + PlanningResult, + PlanningSceneInfo, +) from dimos.manipulation.planning.spec.protocols import VisualizationSpec from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -116,8 +122,7 @@ def __init__(self) -> None: self._error_message = "" self._planning_epoch = 0 self._robots = {} - self._planned_paths = {} - self._planned_trajectories = {} + self._last_plan = None self._world_monitor = None self._planner = None self._kinematics = None @@ -130,6 +135,49 @@ def _make_module() -> ManipulationModule: return _ManipulationModuleHarness() +def _one_joint_config(name: str = "arm") -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" + ) + ], + coordinator_task_name=f"traj_{name}", + ) + + +def _install_generated_plan( + module: ManipulationModule, + config: RobotModelConfig, + traj_gen: MagicMock, + *points: list[float], +) -> None: + """Install a generated plan and enough monitor state to derive robot paths.""" + module._robots = {config.name: ("robot_id", config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=config.joint_names, + position=[0.0 for _ in config.joint_names], + ) + module._last_plan = GeneratedPlan( + group_ids=(f"{config.name}/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=[f"{config.name}/{joint}" for joint in config.joint_names], + position=list(point), + ) + for point in points + ], + ) + + class TestStateMachine: """Test state transitions.""" @@ -323,8 +371,13 @@ def test_solve_ik_rpc_calls_configured_backend(self, robot_config): 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]) current = JointState(name=robot_config.joint_names, position=[0.0, 0.0, 0.0]) - module._world_monitor.get_current_joint_state.return_value = current + current_global = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.current_global_joint_state.return_value = current_global expected = IKResult( status=IKStatus.SUCCESS, joint_state=JointState(name=robot_config.joint_names, position=[0.1, 0.2, 0.3]), @@ -334,29 +387,34 @@ def test_solve_ik_rpc_calls_configured_backend(self, robot_config): message="ok", ) module._kinematics = MagicMock() - module._kinematics.solve.return_value = expected + module._kinematics.solve_pose_targets.return_value = expected pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) result = module.solve_ik(pose) assert result is expected assert module._state == ManipulationState.COMPLETED - assert module._planned_paths == {} - module._kinematics.solve.assert_called_once() - _, kwargs = module._kinematics.solve.call_args + assert module._last_plan is None + module._kinematics.solve_pose_targets.assert_called_once() + _, kwargs = module._kinematics.solve_pose_targets.call_args assert kwargs["world"] is module._world_monitor.world - assert kwargs["robot_id"] == "robot_id" - assert kwargs["seed"] is current + assert kwargs["seed"].name == current_global.name + assert kwargs["seed"].position == current.position assert kwargs["check_collision"] is True - assert kwargs["target_pose"].frame_id == "world" - assert kwargs["target_pose"].position.x == 0.45 + [(group, target_pose)] = kwargs["pose_targets"].items() + assert group.id == "test_arm/manipulator" + assert target_pose.frame_id == "world" + assert target_pose.position.x == 0.45 def test_solve_ik_rpc_returns_failure_without_joint_state(self, robot_config): """solve_ik reports a failed IKResult when no seed state is available.""" module = _make_module() module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} module._world_monitor = MagicMock() - module._world_monitor.get_current_joint_state.return_value = None + module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=[], position=[] + ) module._kinematics = MagicMock() pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) @@ -365,7 +423,7 @@ def test_solve_ik_rpc_returns_failure_without_joint_state(self, robot_config): assert result.status == IKStatus.NO_SOLUTION assert result.message == "No joint state" assert module._state == ManipulationState.IDLE - module._kinematics.solve.assert_not_called() + module._kinematics.solve_pose_targets.assert_not_called() def test_solve_ik_rpc_uses_explicit_seed(self, robot_config): """solve_ik initializes the backend from an explicit seed when provided.""" @@ -373,21 +431,317 @@ def test_solve_ik_rpc_uses_explicit_seed(self, robot_config): module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, position=[0.0, 0.0, 0.0] - ) + module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) explicit_seed = JointState(name=robot_config.joint_names, position=[0.2, 0.1, 0.0]) expected = IKResult(status=IKStatus.SUCCESS, joint_state=explicit_seed) module._kinematics = MagicMock() - module._kinematics.solve.return_value = expected + module._kinematics.solve_pose_targets.return_value = expected pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) result = module.solve_ik(pose, seed=explicit_seed) assert result is expected - _, kwargs = module._kinematics.solve.call_args + _, kwargs = module._kinematics.solve_pose_targets.call_args assert kwargs["seed"] is explicit_seed - module._world_monitor.get_current_joint_state.assert_not_called() + module._world_monitor.current_global_joint_state.assert_not_called() + + +class TestPlanningGroupApis: + """Test explicit planning-group API behavior.""" + + def test_list_planning_groups_and_robot_info_include_groups(self, robot_config): + module = _make_module() + registry = PlanningGroupRegistry([robot_config]) + module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = registry + module._init_joints = {} + + groups = module.list_planning_groups() + info = module.get_robot_info() + + assert [group.id for group in groups] == ["test_arm/manipulator"] + assert info is not None + assert info["planning_groups"] == groups + assert info["end_effector_link"] == "link_tcp" + assert info["has_joint_name_mapping"] is False + + def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches(self, robot_config): + module = _make_module() + registry = PlanningGroupRegistry([robot_config]) + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory( + (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) + ) + module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = registry + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=robot_config.joint_names, + position=[0.0, 0.0, 0.0], + ) + result_path = [ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ), + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ), + ] + module._planner = MagicMock() + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=result_path, + planning_time=0.01, + path_length=0.3, + iterations=4, + message="ok", + ) + + success = module.plan_to_joint_targets( + { + "test_arm/manipulator": JointState( + name=robot_config.joint_names, + position=[0.1, 0.2, 0.3], + ) + } + ) + + assert success is True + assert module._last_plan is not None + assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module._last_plan.path == result_path + projected = module._project_plan_to_robot_paths(module._last_plan) + assert projected is not None + assert projected["test_arm"][1].position == [0.1, 0.2, 0.3] + assert module._trajectory_from_robot_path("test_arm", projected["test_arm"]) is not None + module._planner.plan_selected_joint_path.assert_called_once() + _, kwargs = module._planner.plan_selected_joint_path.call_args + assert kwargs["selection"].group_ids == ("test_arm/manipulator",) + assert kwargs["goal"].name == [ + "test_arm/joint1", + "test_arm/joint2", + "test_arm/joint3", + ] + + def test_plan_to_pose_targets_uses_group_ik_and_selected_path(self, robot_config): + module = _make_module() + registry = PlanningGroupRegistry([robot_config]) + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory( + (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) + ) + module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = registry + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=robot_config.joint_names, + position=[0.0, 0.0, 0.0], + ) + ik_goal = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ) + module._kinematics = MagicMock() + module._kinematics.solve_pose_targets.return_value = IKResult( + status=IKStatus.SUCCESS, + joint_state=ik_goal, + ) + module._planner = MagicMock() + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ), + ik_goal, + ], + ) + pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) + + success = module.plan_to_pose_targets({"test_arm/manipulator": pose}) + + assert success is True + module._kinematics.solve_pose_targets.assert_called_once() + _, ik_kwargs = module._kinematics.solve_pose_targets.call_args + target_groups = list(ik_kwargs["pose_targets"].keys()) + assert [group.id for group in target_groups] == ["test_arm/manipulator"] + target_pose = ik_kwargs["pose_targets"][target_groups[0]] + assert target_pose.position.x == 0.45 + assert ik_kwargs["seed"].name == [ + "test_arm/joint1", + "test_arm/joint2", + "test_arm/joint3", + ] + _, planner_kwargs = module._planner.plan_selected_joint_path.call_args + assert planner_kwargs["goal"] is ik_goal + + def test_failed_plan_projection_clears_generated_plan(self, robot_config, simple_trajectory): + module = _make_module() + registry = PlanningGroupRegistry([robot_config]) + module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = registry + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.get_current_joint_state.return_value = None + module._last_plan = GeneratedPlan( + group_ids=("test_arm/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ), + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.2, 0.2, 0.2], + ), + ], + ) + module._planner = MagicMock() + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ) + ], + ) + + success = module.plan_to_joint_targets( + {"test_arm/manipulator": JointState(position=[0.1, 0.2, 0.3])} + ) + + assert success is False + assert module._state == ManipulationState.FAULT + assert module._last_plan is None + assert module.has_planned_path() is False + + def test_execute_plan_dispatches_each_affected_robot(self, robot_config, simple_trajectory): + module = _make_module() + registry = PlanningGroupRegistry([robot_config]) + traj_gen = MagicMock() + traj_gen.generate.return_value = simple_trajectory + module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = registry + module._world_monitor.get_current_joint_state.return_value = JointState( + name=robot_config.joint_names, + position=[0.0, 0.0, 0.0], + ) + module._last_plan = GeneratedPlan( + group_ids=("test_arm/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ), + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ), + ], + ) + mock_client = MagicMock() + mock_client.task_invoke.return_value = True + module._coordinator_client = mock_client + + assert module.execute_plan() is True + + mock_client.task_invoke.assert_called_once() + task_name, method_name, payload = mock_client.task_invoke.call_args.args + assert task_name == "traj_arm" + assert method_name == "execute" + trajectory = payload["trajectory"] + assert trajectory.joint_names == simple_trajectory.joint_names + assert trajectory.points == simple_trajectory.points + assert module._state == ManipulationState.COMPLETED + + def test_pose_wrappers_fail_safely_without_unique_pose_group(self, robot_config): + no_pose_config = RobotModelConfig( + name="test_arm", + model_path=robot_config.model_path, + base_pose=robot_config.base_pose, + joint_names=robot_config.joint_names, + base_link=robot_config.base_link, + planning_groups=[ + PlanningGroupDefinition( + name="joint_only", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + ) + ], + ) + module = _make_module() + module._robots = {"test_arm": ("robot_id", no_pose_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([no_pose_config]) + module._world_monitor.get_ee_pose.side_effect = ValueError("no pose group") + module._kinematics = MagicMock() + + pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) + + assert module.get_ee_pose() is None + assert module.plan_to_pose(pose) is False + result = module.inverse_kinematics_single(pose) + assert result.status == IKStatus.NO_SOLUTION + assert "no pose-targetable planning group" in result.message + + def test_pose_wrappers_fail_safely_with_multiple_pose_groups(self, robot_config): + multi_pose_config = RobotModelConfig( + name="test_arm", + model_path=robot_config.model_path, + base_pose=robot_config.base_pose, + joint_names=robot_config.joint_names, + base_link=robot_config.base_link, + planning_groups=[ + PlanningGroupDefinition( + name="wrist", + joint_names=("joint1", "joint2"), + base_link="link_base", + tip_link="link_wrist", + ), + PlanningGroupDefinition( + name="tool", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ), + ], + ) + module = _make_module() + module._robots = {"test_arm": ("robot_id", multi_pose_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([multi_pose_config]) + module._world_monitor.get_ee_pose.side_effect = ValueError("multiple pose groups") + module._kinematics = MagicMock() + + pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) + + assert module.get_ee_pose() is None + assert module.plan_to_pose(pose) is False + result = module.inverse_kinematics_single(pose) + assert result.status == IKStatus.NO_SOLUTION + assert "2 pose-targetable planning groups" in result.message class TestJointNameTranslation: @@ -418,9 +772,9 @@ def test_execute_requires_trajectory(self, robot_config): """Execute fails without planned trajectory.""" module = _make_module() module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {} assert module.execute() is False + assert module._state == ManipulationState.IDLE def test_execute_requires_task_name(self): """Execute fails without coordinator_task_name.""" @@ -436,38 +790,87 @@ def test_execute_requires_task_name(self): ) ], ) - module._robots = {"arm": ("id", config_no_task, MagicMock())} - module._planned_trajectories = {"arm": MagicMock()} + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (1.0, [0.1])) + _install_generated_plan(module, config_no_task, traj_gen, [0.0], [0.1]) - assert module.execute() is False + assert module.execute("arm") is False + assert module._state == ManipulationState.IDLE + + def test_execute_plan_requires_task_name_without_sticking_executing(self, robot_config): + module = _make_module() + config_no_task = RobotModelConfig( + name="test_arm", + model_path=robot_config.model_path, + base_pose=robot_config.base_pose, + joint_names=robot_config.joint_names, + base_link=robot_config.base_link, + planning_groups=robot_config.planning_groups, + gripper_hardware_id="test_gripper", + ) + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory( + (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) + ) + module._robots = {"test_arm": ("id", config_no_task, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config_no_task]) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=robot_config.joint_names, + position=[0.0, 0.0, 0.0], + ) + module._coordinator_client = MagicMock() + module._last_plan = GeneratedPlan( + group_ids=("test_arm/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ), + JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ), + ], + ) + + assert module.execute_plan() is False + assert module._state == ManipulationState.FAULT def test_execute_success(self, robot_config, simple_trajectory): """Successful execute calls coordinator via task_invoke.""" module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {"test_arm": simple_trajectory} + traj_gen = MagicMock() + traj_gen.generate.return_value = simple_trajectory + _install_generated_plan(module, robot_config, traj_gen, [0.0, 0.0, 0.0], [0.1, 0.2, 0.3]) mock_client = MagicMock() mock_client.task_invoke.return_value = True module._coordinator_client = mock_client - assert module.execute() is True + assert module.execute("test_arm") is True assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_called_once_with( - "traj_arm", "execute", {"trajectory": simple_trajectory} - ) + mock_client.task_invoke.assert_called_once() + task_name, method_name, payload = mock_client.task_invoke.call_args.args + assert task_name == "traj_arm" + assert method_name == "execute" + trajectory = payload["trajectory"] + assert trajectory.joint_names == simple_trajectory.joint_names + assert trajectory.points == simple_trajectory.points def test_execute_rejected(self, robot_config, simple_trajectory): """Rejected execution sets FAULT state.""" module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {"test_arm": simple_trajectory} + traj_gen = MagicMock() + traj_gen.generate.return_value = simple_trajectory + _install_generated_plan(module, robot_config, traj_gen, [0.0, 0.0, 0.0], [0.1, 0.2, 0.3]) mock_client = MagicMock() mock_client.task_invoke.return_value = False module._coordinator_client = mock_client - assert module.execute() is False + assert module.execute("test_arm") is False assert module._state == ManipulationState.FAULT @@ -733,10 +1136,10 @@ def test_dismiss_preview_routes_to_monitor(self): def test_preview_path_uses_trajectory_duration_and_interpolates(self): module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [2.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (2.0, [2.0]))} + config = _one_joint_config() + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (2.0, [2.0])) + _install_generated_plan(module, config, traj_gen, [0.0], [2.0]) assert module.preview_path(robot_name="arm", target_fps=2.0) is True @@ -748,10 +1151,10 @@ def test_preview_path_uses_trajectory_duration_and_interpolates(self): def test_preview_path_explicit_duration_overrides_and_fps_densifies(self): module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [9.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (9.0, [9.0]))} + config = _one_joint_config() + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (9.0, [9.0])) + _install_generated_plan(module, config, traj_gen, [0.0], [9.0]) assert module.preview_path(duration=1.5, robot_name="arm", target_fps=2.0) is True @@ -761,42 +1164,39 @@ def test_preview_path_explicit_duration_overrides_and_fps_densifies(self): assert duration == 1.5 assert [state.position for state in preview_path] == [[0.0], [3.0], [6.0], [9.0]] - def test_preview_path_missing_trajectory_uses_default_duration(self): + def test_preview_path_returns_false_when_trajectory_cannot_be_derived(self): module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._planned_trajectories = {} + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0]) - assert module.preview_path(robot_name="arm", target_fps=10.0) is True - - module._world_monitor.animate_path.assert_called_once_with( - "robot_id", module._planned_paths["arm"], 3.0 - ) + assert module.preview_path(robot_name="arm", target_fps=10.0) is False + module._world_monitor.animate_path.assert_not_called() def test_preview_path_skips_interpolation_for_nonpositive_fps_or_duration(self): module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (2.0, [1.0]))} + config = _one_joint_config() + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (2.0, [1.0])) + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) assert module.preview_path(robot_name="arm", target_fps=0.0) is True assert module.preview_path(duration=0.0, robot_name="arm", target_fps=20.0) is True - assert ( - module._world_monitor.animate_path.call_args_list[0].args[1] - == module._planned_paths["arm"] - ) - assert ( - module._world_monitor.animate_path.call_args_list[1].args[1] - == module._planned_paths["arm"] - ) + assert [ + state.position for state in module._world_monitor.animate_path.call_args_list[0].args[1] + ] == [[0.0], [1.0]] + assert [ + state.position for state in module._world_monitor.animate_path.call_args_list[1].args[1] + ] == [[0.0], [1.0]] def test_preview_path_returns_false_for_missing_inputs(self): module = _make_module() - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} + config = _one_joint_config() + traj_gen = MagicMock() + traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (1.0, [1.0])) + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) + module._world_monitor = None assert module.preview_path(robot_name="arm") is False @@ -804,6 +1204,5 @@ def test_preview_path_returns_false_for_missing_inputs(self): module._robots = {} assert module.preview_path(robot_name="arm") is False - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": []} + _install_generated_plan(module, config, traj_gen) assert module.preview_path(robot_name="arm") is False diff --git a/dimos/manipulation/visualization/viser/adapter.py b/dimos/manipulation/visualization/viser/adapter.py index c3eb8e360b..a3f363d8e4 100644 --- a/dimos/manipulation/visualization/viser/adapter.py +++ b/dimos/manipulation/visualization/viser/adapter.py @@ -24,7 +24,7 @@ from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor from dimos.manipulation.planning.spec.config import RobotModelConfig - from dimos.manipulation.planning.spec.models import JointPath, RobotName, WorldRobotID + from dimos.manipulation.planning.spec.models import RobotName, WorldRobotID from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -34,6 +34,18 @@ def copy_joint_state(joint_state: JointState | None) -> JointState | None: return None if joint_state is None else JointState(joint_state) +def _string_list(value: object) -> list[str]: + return [str(item) for item in value] if isinstance(value, list) else [] + + +def _float_list_or_none(value: object) -> list[float] | None: + return [float(item) for item in value] if isinstance(value, list) else None + + +def _float_value(value: object, default: float = 0.0) -> float: + return float(value) if isinstance(value, int | float | str) else default + + class InProcessViserAdapter: """Small in-process boundary between Viser callbacks and manipulation internals.""" @@ -65,26 +77,24 @@ def get_robot_info(self, robot_name: RobotName) -> RobotInfo | None: info = self._module.get_robot_info(robot_name) if info is None: return None - return { + + robot_info: RobotInfo = { "name": str(info["name"]), "world_robot_id": str(info["world_robot_id"]), - "joint_names": [str(name) for name in info["joint_names"]], + "joint_names": _string_list(info["joint_names"]), "end_effector_link": str(info["end_effector_link"]), "base_link": str(info["base_link"]), - "max_velocity": float(info["max_velocity"]), - "max_acceleration": float(info["max_acceleration"]), + "max_velocity": _float_value(info["max_velocity"]), + "max_acceleration": _float_value(info["max_acceleration"]), "has_joint_name_mapping": bool(info["has_joint_name_mapping"]), "coordinator_task_name": None if info["coordinator_task_name"] is None else str(info["coordinator_task_name"]), - "home_joints": None - if info["home_joints"] is None - else [float(value) for value in info["home_joints"]], - "pre_grasp_offset": float(info["pre_grasp_offset"]), - "init_joints": None - if info["init_joints"] is None - else [float(value) for value in info["init_joints"]], + "home_joints": _float_list_or_none(info["home_joints"]), + "pre_grasp_offset": _float_value(info["pre_grasp_offset"]), + "init_joints": _float_list_or_none(info["init_joints"]), } + return robot_info def get_init_joints(self, robot_name: RobotName) -> JointState | None: return copy_joint_state(self._module.get_init_joints(robot_name)) @@ -127,16 +137,6 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu ) return result - def get_planned_path(self, robot_name: RobotName) -> JointPath | None: - path = self._module.get_planned_path(robot_name) - if path is None: - return None - copied = [copy_joint_state(point) for point in path] - return [point for point in copied if point is not None] - - def get_planned_trajectory_duration(self, robot_name: RobotName) -> float | None: - return self._module.get_planned_trajectory_duration(robot_name) - def get_module_state(self) -> str: return str(self._module.get_state()) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index 2c0abac52d..218290c5c4 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -579,13 +579,11 @@ def operation() -> None: if not self._operation_is_current(operation_id): return if ok: - path = self.adapter.get_planned_path(robot_name) self.state.plan_state.status = PlanStatus.FRESH self.state.plan_state.robot = robot_name self.state.plan_state.target_joints = list(target.position) self.state.plan_state.target_pose = self.state.cartesian_target self.state.plan_state.start_joints_snapshot = list(self.state.current_joints or []) - self.state.plan_state.planned_path = path else: self.state.plan_state.status = PlanStatus.FAILED self._finish_operation(f"plan_to_joints={ok}", operation_id=operation_id) diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index b46097df95..d8500692fd 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -96,7 +96,6 @@ class PanelPlanState: target_pose: Pose | None = None target_joints: list[float] | None = None start_joints_snapshot: list[float] | None = None - planned_path: list[JointState] | None = None @dataclass diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 2994158022..09976b05f8 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -396,13 +396,6 @@ def get_robot_info(self, robot_name: str) -> RobotInfo | None: def get_init_joints(self, robot_name: str) -> JointState | None: return getattr(self, "_init_joints", {}).get(robot_name) - def get_planned_path(self, robot_name: str) -> list[JointState] | None: - return getattr(self, "_planned_paths", {}).get(robot_name) - - def get_planned_trajectory_duration(self, robot_name: str) -> float | None: - trajectory = getattr(self, "_planned_trajectories", {}).get(robot_name) - return None if trajectory is None else float(trajectory.duration) - def get_state(self) -> str: state = getattr(self, "_state", "IDLE") return str(getattr(state, "name", state)) @@ -455,8 +448,6 @@ def make_adapter_with_robot() -> InProcessViserAdapter: module = FakeManipulationModule( _robots={"arm": ("robot-1", config, None)}, _init_joints={"arm": FakeJointState(["j1", "j2"], position=[0.1, 0.2])}, - _planned_paths={}, - _planned_trajectories={}, _state=NamedState(name="IDLE"), _error_message="", _world_monitor=world_monitor, @@ -903,8 +894,6 @@ def test_adapter_copies_joint_state_and_delegates_to_module() -> None: copied = FakeJointState(["j1"], position=[1.0], velocity=[2.0], effort=[3.0]) module = FakeManipulationModule( _robots={"arm": ("robot-1", SimpleNamespace(), None)}, - _planned_paths={"arm": [copied]}, - _planned_trajectories={}, plan_to_pose=lambda pose, robot_name=None: (pose, robot_name), plan_to_joints=lambda joints, robot_name=None: (joints, robot_name), preview_path=lambda robot_name=None: robot_name, @@ -921,19 +910,13 @@ def test_adapter_copies_joint_state_and_delegates_to_module() -> None: module._world_monitor = world_monitor adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - planned = adapter.get_planned_path("arm") - assert planned is not None - assert planned[0] is not copied - assert planned[0].name is not copied.name - assert planned[0].position is not copied.position - current = adapter.get_current_joint_state("arm") assert current is not copied assert current.name is not copied.name assert adapter.plan_to_pose("pose", "arm") == ("pose", "arm") assert adapter.preview_path("arm") == "arm" - assert adapter.evaluate_joint_target(planned[0], "arm")["status"] == "FEASIBLE" + assert adapter.evaluate_joint_target(copied, "arm")["status"] == "FEASIBLE" def test_adapter_evaluate_joint_target_uses_world_monitor_and_copies_input() -> None: @@ -1112,9 +1095,7 @@ def test_gui_initializes_pose_selector_to_current_ee_pose( orientation=SimpleNamespace(w=0.9, x=0.1, y=0.2, z=0.3), ) config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) + module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, is_state_stale=lambda robot_id, max_age=1.0: False, @@ -1139,8 +1120,6 @@ def test_gui_preset_dropdown_and_controls_include_init_home_current_and_callback module = FakeManipulationModule( _robots={"arm": ("robot-1", config, None)}, _init_joints={"arm": FakeJointState(["j1", "j2"], position=[-1.0, -2.0])}, - _planned_paths={}, - _planned_trajectories={}, ) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, @@ -1165,9 +1144,7 @@ def test_gui_rebuilding_joint_sliders_removes_stale_viser_handles( ) -> None: current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) config = make_robot_config(joint_names=["j1", "j2"], home_joints=[1.0, 2.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) + module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, is_state_stale=lambda robot_id, max_age=1.0: False, @@ -1212,8 +1189,6 @@ def test_panel_execution_requires_fresh_plan_and_refresh_updates_robot_controls( config = make_robot_config(joint_names=["j1"], home_joints=[0.5]) module = FakeManipulationModule( _robots={"arm": ("robot-1", config, None)}, - _planned_paths={}, - _planned_trajectories={}, execute=lambda robot_name=None: False, ) world_monitor = SimpleNamespace( @@ -1243,9 +1218,7 @@ def test_gui_moves_joint_target_immediately_and_stores_evaluated_joint_solution( current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) target_pose = SimpleNamespace(position=SimpleNamespace(x=0.2, y=0.3, z=0.4)) config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) + module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, is_state_stale=lambda robot_id, max_age=1.0: False, @@ -1308,9 +1281,7 @@ def test_gui_cartesian_ik_result_does_not_rewrite_active_gizmo( ) -> None: current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) + module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, is_state_stale=lambda robot_id, max_age=1.0: False, @@ -1354,9 +1325,7 @@ def test_gui_collision_evaluation_marks_target_infeasible_and_colors_scene( ) -> None: current = FakeJointState(["j1"], position=[0.0]) config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) + module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) world_monitor = SimpleNamespace( get_current_joint_state=lambda robot_id: current, is_state_stale=lambda robot_id, max_age=1.0: False, @@ -1393,15 +1362,12 @@ def test_gui_safe_execute_requires_fresh_matching_plan_and_clear_resets_path( make_panel: Callable[..., ViserPanelGui], monkeypatch: pytest.MonkeyPatch ) -> None: current = FakeJointState(["j1"], position=[1.0]) - planned = [FakeJointState(["j1"], position=[1.0]), FakeJointState(["j1"], position=[2.0])] executed = [] cleared = [] module = FakeManipulationModule( _robots={ "arm": ("robot-1", make_robot_config(joint_names=["j1"], home_joints=[1.0]), None) }, - _planned_paths={"arm": planned}, - _planned_trajectories={}, _state=NamedState(name="IDLE"), execute=lambda robot_name=None: executed.append(robot_name) or True, clear_planned_path=lambda: cleared.append(True) or True, @@ -1433,7 +1399,6 @@ def test_gui_safe_execute_requires_fresh_matching_plan_and_clear_resets_path( status=PlanStatus.FRESH, robot="arm", start_joints_snapshot=[1.2], - planned_path=planned, ) gui._submit_execute() assert executed == [] diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..6b4dd8c68e 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -60,6 +60,7 @@ "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic", "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic", "dual-xarm6-planner": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner", + "dual-xarm6-planner-coordinator-mock-meshcat": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator_mock_meshcat", "keyboard-teleop-a750": "dimos.robot.manipulators.a750.blueprints.teleop:keyboard_teleop_a750", "keyboard-teleop-openarm": "dimos.robot.manipulators.openarm.blueprints.teleop:keyboard_teleop_openarm", "keyboard-teleop-openarm-mock": "dimos.robot.manipulators.openarm.blueprints.teleop:keyboard_teleop_openarm_mock", diff --git a/dimos/robot/manipulators/test_blueprints.py b/dimos/robot/manipulators/test_blueprints.py index 765926f8ac..c8bd7ddaa9 100644 --- a/dimos/robot/manipulators/test_blueprints.py +++ b/dimos/robot/manipulators/test_blueprints.py @@ -14,12 +14,17 @@ from typing import Any +from dimos.control.coordinator import ControlCoordinator from dimos.core.coordination.blueprints import Blueprint from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationModuleConfig -from dimos.manipulation.visualization.config import NoManipulationVisualizationConfig +from dimos.manipulation.visualization.config import ( + MeshcatVisualizationConfig, + NoManipulationVisualizationConfig, +) from dimos.robot.manipulators.common.blueprints import planner from dimos.robot.manipulators.xarm.blueprints.basic import ( dual_xarm6_planner, + dual_xarm6_planner_coordinator_mock_meshcat, xarm6_planner_only, xarm7_planner_coordinator, ) @@ -58,3 +63,23 @@ def test_xarm_planner_blueprints_default_to_no_visualization() -> None: config = _manipulation_config(blueprint) assert isinstance(config.visualization, NoManipulationVisualizationConfig) + + +def test_dual_xarm6_mock_meshcat_blueprint_wires_planner_and_coordinator() -> None: + config = _manipulation_config(dual_xarm6_planner_coordinator_mock_meshcat) + coordinator_kwargs = next( + atom.kwargs + for atom in dual_xarm6_planner_coordinator_mock_meshcat.blueprints + if atom.module is ControlCoordinator + ) + + assert isinstance(config.visualization, MeshcatVisualizationConfig) + assert [robot.name for robot in config.robots] == ["left_arm", "right_arm"] + assert [hardware.hardware_id for hardware in coordinator_kwargs["hardware"]] == [ + "left_arm", + "right_arm", + ] + assert [task.name for task in coordinator_kwargs["tasks"]] == [ + "traj_left_arm", + "traj_right_arm", + ] diff --git a/dimos/robot/manipulators/xarm/blueprints/basic.py b/dimos/robot/manipulators/xarm/blueprints/basic.py index f66d50079b..68e6bb467e 100644 --- a/dimos/robot/manipulators/xarm/blueprints/basic.py +++ b/dimos/robot/manipulators/xarm/blueprints/basic.py @@ -26,6 +26,7 @@ XARM7_SIM_PATH, make_xarm6_model_config, make_xarm7_model_config, + make_xarm_hardware, xarm6_hardware, xarm7_hardware, ) @@ -43,6 +44,26 @@ planning_timeout=10.0, ) +_mock_left_xarm6_hw = make_xarm_hardware("left_arm", 6) +_mock_right_xarm6_hw = make_xarm_hardware("right_arm", 6) + +dual_xarm6_planner_coordinator_mock_meshcat = autoconnect( + planner( + robots=[ + make_xarm6_model_config(name="left_arm", y_offset=0.5), + make_xarm6_model_config(name="right_arm", y_offset=-0.5), + ], + visualization={"backend": "meshcat"}, + ), + coordinator( + hardware=[_mock_left_xarm6_hw, _mock_right_xarm6_hw], + tasks=[ + trajectory_task(_mock_left_xarm6_hw), + trajectory_task(_mock_right_xarm6_hw), + ], + ), +) + _xarm7_hw = xarm7_hardware("arm", gripper=True, mock_without_address=True) xarm7_planner_coordinator = autoconnect( diff --git a/dimos/robot/test_all_blueprints.py b/dimos/robot/test_all_blueprints.py index ec51e681f5..ba907315fe 100644 --- a/dimos/robot/test_all_blueprints.py +++ b/dimos/robot/test_all_blueprints.py @@ -50,6 +50,7 @@ "coordinator-xarm6", "coordinator-xarm7", "dual-xarm6-planner", + "dual-xarm6-planner-coordinator-mock-meshcat", "learning-collect-quest-piper", "learning-collect-quest-xarm7", "teleop-hosted-go2", diff --git a/openspec/changes/manipulation-module-group-api/.openspec.yaml b/openspec/changes/manipulation-module-group-api/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/manipulation-module-group-api/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/manipulation-module-group-api/README.md b/openspec/changes/manipulation-module-group-api/README.md deleted file mode 100644 index a4cb1c316f..0000000000 --- a/openspec/changes/manipulation-module-group-api/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# manipulation-module-group-api - -Extract PR 4 from reference branch cc/spec/movegroup: expose group-aware planning through ManipulationModule while preserving safe robot-scoped compatibility behavior. diff --git a/openspec/changes/manipulation-module-group-api/design.md b/openspec/changes/manipulation-module-group-api/design.md deleted file mode 100644 index 5e84d72704..0000000000 --- a/openspec/changes/manipulation-module-group-api/design.md +++ /dev/null @@ -1,28 +0,0 @@ -## Context - -The manipulation module is where planning groups become user-facing. It must expose explicit group APIs while retaining safe behavior for existing robot-scoped wrappers. - -## Goals / Non-Goals - -**Goals:** -- Support explicit group IDs for joint targets, pose targets, IK, previews, and robot info. -- Keep robot-scoped wrappers deterministic: no unique pose group means no successful robot-scoped pose action. -- Preserve safe return contracts for RPC-like public methods. -- Cover compatibility behavior in unit tests. - -**Non-Goals:** -- Do not implement Viser panel behavior in this PR. -- Do not add new planner algorithms. -- Do not change control task behavior unless required by module API compilation. - -## Decisions - -- `get_ee_pose` returns `None` when no unique pose-targetable group exists. -- `plan_to_pose` returns `False` when no unique pose-targetable group exists. -- `inverse_kinematics_single` returns `IKResult(NO_SOLUTION, message=...)` when no unique pose-targetable group exists. -- Explicit group APIs are preferred for multi-group robots. - -## Risks / Trade-offs - -- This file has a large diff. Keep the PR focused on API behavior and tests; do not include UI state-machine changes. -- Some compatibility wrappers are intentionally conservative to avoid silently selecting the wrong group. diff --git a/openspec/changes/manipulation-module-group-api/proposal.md b/openspec/changes/manipulation-module-group-api/proposal.md deleted file mode 100644 index a74f5a3adf..0000000000 --- a/openspec/changes/manipulation-module-group-api/proposal.md +++ /dev/null @@ -1,25 +0,0 @@ -## Why - -Once engines and algorithms are group-aware, the public manipulation module must expose those capabilities safely. This PR is the reviewer-facing API migration and should stay separate from the Viser UI rewrite. - -## What Changes - -- Add or migrate `ManipulationModule` APIs for group-aware joint targets, pose targets, previews, IK, and robot info. -- Preserve robot-scoped compatibility wrappers with explicit behavior when there is no unique pose-targetable group. -- Update coordinator client and example client call sites. -- Add module/unit/e2e tests for public behavior. - -## Capabilities - -### New Capabilities -- `manipulation-module-group-api`: Public manipulation APIs support planning-group IDs and safe compatibility wrappers. - -### Modified Capabilities -- `manipulation-module-planning`: Robot-scoped pose wrappers no longer assume an arbitrary or hidden end-effector. - -## Impact - -- Base branch: PR 3 `group-aware-ik-rrt`. -- Reference implementation: `cc/spec/movegroup`, including Greptile follow-up behavior for robot-scoped wrappers. -- Primary files: `dimos/manipulation/manipulation_module.py`, coordinator/example clients, manipulation module tests, and planning-group e2e test. -- Out of scope: Viser UI and control/task changes. diff --git a/openspec/changes/manipulation-module-group-api/tasks.md b/openspec/changes/manipulation-module-group-api/tasks.md deleted file mode 100644 index 7f50e78293..0000000000 --- a/openspec/changes/manipulation-module-group-api/tasks.md +++ /dev/null @@ -1,18 +0,0 @@ -## 1. Extract module API changes - -- [ ] 1.1 Start from PR 3 and use `cc/spec/movegroup` as reference. -- [ ] 1.2 Extract group-aware target, IK, preview, robot-info, and execution behavior from `ManipulationModule`. -- [ ] 1.3 Extract coordinator client and example client updates only as needed for the public API. -- [ ] 1.4 Include Greptile follow-up behavior for `get_ee_pose` and `plan_to_pose`. - -## 2. Tests - -- [ ] 2.1 Bring over focused `test_manipulation_unit.py` coverage for group APIs and compatibility wrappers. -- [ ] 2.2 Bring over relevant `test_manipulation_module.py` integration coverage. -- [ ] 2.3 Bring over `dimos/e2e_tests/test_manipulation_planning_groups.py` if it can run against this stacked base. - -## 3. Validation - -- [ ] 3.1 Run `uv run pytest dimos/manipulation/test_manipulation_unit.py dimos/manipulation/test_manipulation_module.py dimos/e2e_tests/test_manipulation_planning_groups.py -q`. -- [ ] 3.2 Run `uv run mypy dimos/manipulation/manipulation_module.py`. -- [ ] 3.3 Verify no Viser implementation files are included. diff --git a/openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md b/openspec/specs/manipulation-module-group-api/spec.md similarity index 88% rename from openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md rename to openspec/specs/manipulation-module-group-api/spec.md index 16aa9bb6bc..b812df591f 100644 --- a/openspec/changes/manipulation-module-group-api/specs/manipulation-module-group-api/spec.md +++ b/openspec/specs/manipulation-module-group-api/spec.md @@ -1,4 +1,8 @@ -## ADDED Requirements +## Purpose + +Define the public manipulation module APIs for explicit planning-group targeting and safe robot-scoped compatibility behavior. + +## Requirements ### Requirement: Manipulation module must expose explicit group planning APIs The manipulation module MUST allow callers to target planning groups explicitly for joint and pose planning operations. From 89c3a53ad3513e1b703c8b718737a0d0a870b97f Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:51:59 -0700 Subject: [PATCH 19/29] Planning Group 5: Add group-aware Viser planning (#2860) --- CONTEXT.md | 53 + .../tasks/trajectory_task/trajectory_task.py | 72 +- dimos/control/test_control.py | 135 + .../_test_manipulation_helpers.py | 45 + dimos/manipulation/manipulation_module.py | 673 +++-- .../planning/monitor/test_world_monitor.py | 60 +- .../planning/monitor/world_monitor.py | 96 +- dimos/manipulation/planning/spec/models.py | 24 +- dimos/manipulation/planning/spec/protocols.py | 27 +- .../planning/world/drake_world.py | 172 +- .../world/test_drake_world_planning_groups.py | 189 ++ .../test_generated_plan_materialization.py | 181 ++ .../manipulation/test_manipulation_module.py | 16 +- .../test_manipulation_monitor_preview.py | 445 ++++ dimos/manipulation/test_manipulation_unit.py | 622 ++--- .../test_plan_execution_reservation.py | 271 ++ dimos/manipulation/visualization/factory.py | 7 +- dimos/manipulation/visualization/operator.py | 343 +++ .../visualization/test_factory.py | 57 +- .../visualization/test_operator.py | 390 +++ .../visualization/viser/adapter.py | 174 -- .../visualization/viser/animation.py | 106 +- dimos/manipulation/visualization/viser/gui.py | 1110 ++++++-- .../manipulation/visualization/viser/scene.py | 323 ++- .../manipulation/visualization/viser/state.py | 63 +- .../visualization/viser/test_gui.py | 474 ++++ .../visualization/viser/test_gui_status.py | 63 - .../viser/test_operation_worker.py | 322 --- .../visualization/viser/test_state.py | 92 +- .../viser/test_viser_visualization.py | 2274 +++++++---------- .../viser/test_visualizer_lifecycle.py | 173 +- .../manipulation/visualization/viser/theme.py | 6 +- .../visualization/viser/visualizer.py | 167 +- dimos/robot/all_blueprints.py | 1 + dimos/robot/manipulators/test_blueprints.py | 19 +- .../manipulators/xarm/blueprints/basic.py | 17 + dimos/robot/test_all_blueprints.py | 1 + .../manipulation/planning_groups.md | 21 +- docs/capabilities/manipulation/readme.md | 40 +- .../.openspec.yaml | 2 + .../design.md | 125 + .../proposal.md | 33 + .../manipulation-module-group-api/spec.md | 56 + .../specs/manipulation-operator/spec.md | 73 + .../spec.md | 64 + .../spec.md | 45 + .../tasks.md | 48 + .../changes/viser-group-planning-ui/design.md | 12 + .../viser-group-planning-ui/proposal.md | 7 +- .../specs/viser-group-planning-ui/spec.md | 73 +- .../changes/viser-group-planning-ui/tasks.md | 18 +- pyproject.toml | 2 + 52 files changed, 6728 insertions(+), 3154 deletions(-) create mode 100644 CONTEXT.md create mode 100644 dimos/manipulation/_test_manipulation_helpers.py create mode 100644 dimos/manipulation/test_generated_plan_materialization.py create mode 100644 dimos/manipulation/test_manipulation_monitor_preview.py create mode 100644 dimos/manipulation/test_plan_execution_reservation.py create mode 100644 dimos/manipulation/visualization/operator.py create mode 100644 dimos/manipulation/visualization/test_operator.py delete mode 100644 dimos/manipulation/visualization/viser/adapter.py create mode 100644 dimos/manipulation/visualization/viser/test_gui.py delete mode 100644 dimos/manipulation/visualization/viser/test_gui_status.py delete mode 100644 dimos/manipulation/visualization/viser/test_operation_worker.py create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/design.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/proposal.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md create mode 100644 openspec/changes/clean-manipulation-visualization-boundary/tasks.md diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 0000000000..2f418448a6 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,53 @@ +# DimOS Robotics + +DimOS composes robot capabilities into plans that can be previewed and executed across one or more robots. + +## Language + +**Planning**: +The complete process of finding a geometric motion and parameterizing it in time for preview and execution. +_Avoid_: Using “planning” to mean geometric path finding alone + +**Path**: +A geometric motion through configuration space without assigned timing, represented by planner-produced waypoints. +_Avoid_: Trajectory + +**Waypoint**: +A discrete robot configuration emitted by a planner as part of a path and consumed by trajectory parameterization. +_Avoid_: Trajectory point, path point + +**Planned Joint**: +A joint explicitly included in a generated plan’s waypoints and synchronized trajectory. A generated plan does not command joints outside this set. +_Avoid_: Every joint belonging to an affected robot + +**Generated Plan**: +The completed outcome of planning, containing both its geometric waypoints and one synchronized parameterized motion. An incomplete plan may exist transiently during construction but is not externally observable. +_Avoid_: Path, planner result + +**Plan Freshness**: +The condition that current planned-joint positions still match a generated plan’s synchronized trajectory start within the accepted tolerance. +_Avoid_: UI snapshot freshness, telemetry freshness + +**Synchronized Trajectory**: +A parameterized motion on one shared relative clock for every planned joint. +_Avoid_: Per-robot trajectory collection + +**Trajectory Parameterization**: +The assignment of timing and motion derivatives to a geometric waypoint path, producing the motion used by preview and execution. +_Avoid_: Preview sampling, path projection + +**Manipulation Operator**: +The human-facing control boundary attached to a manipulation visualization session, through which a client inspects manipulation state and requests target evaluation, planning, preview, execution, cancellation, or reset. It is shared by visualization clients rather than belonging to a particular renderer. +_Avoid_: Visualization backend, renderer + +**Target Draft**: +A frontend-owned, uncommitted set of desired joint or pose targets. A target draft is not a generated plan and has no execution authority. +_Avoid_: Plan, planned path + +**Target Evaluation**: +An advisory assessment of a target draft’s feasibility and resulting joint or pose state. It is feedback for an operator and is not planning or execution authority. +_Avoid_: Generated plan, validated plan + +**Preview**: +A transient visual playback of a generated plan’s synchronized trajectory. A preview has no execution authority and disappears when playback finishes or is cancelled. +_Avoid_: Target ghost, generated plan diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index a3eb23dec3..e3deca4f11 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -23,6 +23,7 @@ from __future__ import annotations from dataclasses import dataclass +import math from typing import Any from dimos.control.task import ( @@ -127,7 +128,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: Returns: JointCommandOutput with positions, or None if not executing """ - if self._trajectory is None: + if self._trajectory is None or not self._trajectory.joint_names: return None # Set start time on first compute() for consistent timing @@ -143,8 +144,10 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: logger.info(f"Trajectory {self._name} completed after {t_elapsed:.3f}s") # Return final position to hold at goal q_ref, _ = self._trajectory.sample(self._trajectory.duration) + final_names = list(self._trajectory.joint_names) + self._clear_active_trajectory() return JointCommandOutput( - joint_names=self._joint_names_list, + joint_names=final_names, positions=list(q_ref), mode=ControlMode.SERVO_POSITION, ) @@ -153,7 +156,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: q_ref, _ = self._trajectory.sample(t_elapsed) return JointCommandOutput( - joint_names=self._joint_names_list, + joint_names=list(self._trajectory.joint_names), positions=list(q_ref), mode=ControlMode.SERVO_POSITION, ) @@ -169,6 +172,60 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: # Abort if any of our joints were preempted if joints & self._joint_names: self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() + + def _clear_active_trajectory(self) -> None: + """Clear stored trajectory-specific execution state.""" + self._trajectory = None + self._pending_start = False + self._start_time = 0.0 + + def _validate_trajectory(self, trajectory: JointTrajectory) -> bool: + """Validate a trajectory before execution.""" + joint_names = list(trajectory.joint_names) + if not joint_names: + logger.warning("Trajectory for %s has empty joint names", self._name) + return False + if len(set(joint_names)) != len(joint_names): + logger.warning("Trajectory for %s has duplicate joint names", self._name) + return False + unknown = [name for name in joint_names if name not in self._joint_names] + if unknown: + logger.warning("Trajectory for %s has unknown joints: %s", self._name, unknown) + return False + if not trajectory.points: + logger.warning("Empty trajectory for %s", self._name) + return False + width = len(joint_names) + previous_time: float | None = None + for index, point in enumerate(trajectory.points): + if len(point.positions) != width or len(point.velocities) != width: + logger.warning("Trajectory point %d for %s has invalid width", index, self._name) + return False + if not all(math.isfinite(value) for value in point.positions): + logger.warning( + "Trajectory point %d for %s has non-finite positions", index, self._name + ) + return False + if not all(math.isfinite(value) for value in point.velocities): + logger.warning( + "Trajectory point %d for %s has non-finite velocities", index, self._name + ) + return False + if not math.isfinite(point.time_from_start): + logger.warning("Trajectory point %d for %s has non-finite time", index, self._name) + return False + if index == 0 and point.time_from_start != 0.0: + logger.warning("Trajectory for %s must start at t=0", self._name) + return False + if previous_time is not None and point.time_from_start <= previous_time: + logger.warning("Trajectory for %s has non-increasing timestamps", self._name) + return False + previous_time = point.time_from_start + if trajectory.duration <= 0.0: + logger.warning("Trajectory for %s has nonpositive duration", self._name) + return False + return True def execute(self, trajectory: JointTrajectory) -> bool: """Start executing a trajectory. @@ -183,17 +240,17 @@ def execute(self, trajectory: JointTrajectory) -> bool: logger.warning(f"Cannot execute: {self._name} in FAULT state") return False - if trajectory is None or trajectory.duration <= 0: + if trajectory is None: logger.warning(f"Invalid trajectory for {self._name}") return False - if not trajectory.points: - logger.warning(f"Empty trajectory for {self._name}") + if not self._validate_trajectory(trajectory): return False # Preempt any active trajectory if self._state == TrajectoryState.EXECUTING: logger.info(f"Preempting active trajectory on {self._name}") + self._clear_active_trajectory() self._trajectory = trajectory self._pending_start = True # Start time set on first compute() @@ -214,6 +271,7 @@ def cancel(self) -> bool: if self._state != TrajectoryState.EXECUTING: return False self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() logger.info(f"Trajectory {self._name} cancelled") return True @@ -227,7 +285,7 @@ def reset(self) -> bool: logger.warning(f"Cannot reset {self._name} while executing") return False self._state = TrajectoryState.IDLE - self._trajectory = None + self._clear_active_trajectory() logger.info(f"Trajectory {self._name} reset to IDLE") return True diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index ae6bc1e9de..443441fad4 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -274,6 +274,141 @@ def test_execute_trajectory(self, trajectory_task, simple_trajectory): assert trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.EXECUTING + def test_execute_partial_subset_and_claims_full_configuration(self, trajectory_task): + trajectory = JointTrajectory( + joint_names=["arm/joint2", "arm/joint3"], + points=[ + TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0], time_from_start=0.0), + TrajectoryPoint(positions=[0.5, 1.0], velocities=[0.0, 0.0], time_from_start=1.0), + ], + ) + + assert trajectory_task.execute(trajectory) is True + assert trajectory_task.claim().joints == frozenset( + {"arm/joint1", "arm/joint2", "arm/joint3"} + ) + + @pytest.mark.parametrize( + "trajectory", + [ + JointTrajectory( + joint_names=[], + points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[])], + ), + JointTrajectory( + joint_names=["arm/joint1", "arm/joint1"], + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0, 0.0], velocities=[0.0, 0.0] + ) + ], + ), + JointTrajectory( + joint_names=["arm/missing"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory(joint_names=["arm/joint1"], points=[]), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], + points=[ + TrajectoryPoint( + time_from_start=0.0, + positions=[0.0, 0.0, 0.0], + velocities=[0.0, 0.0, 0.0], + ) + ], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[float("nan")], velocities=[0.0]) + ], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.1, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), + TrajectoryPoint(time_from_start=0.0, positions=[1.0], velocities=[0.0]), + ], + ), + ], + ) + def test_invalid_partial_inputs_reject_before_state_changes(self, trajectory_task, trajectory): + assert trajectory_task.get_state() == TrajectoryState.IDLE + assert trajectory_task.execute(trajectory) is False + assert trajectory_task.get_state() == TrajectoryState.IDLE + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=0.0, dt=0.01)) + is None + ) + + def test_compute_emits_active_subset_only_and_clears_on_completion(self, trajectory_task): + trajectory = JointTrajectory( + joint_names=["arm/joint2"], + points=[ + TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + ], + ) + assert trajectory_task.execute(trajectory) is True + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.0, dt=0.01)) + output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.5, dt=0.01)) + assert output is not None + assert output.joint_names == ["arm/joint2"] + assert output.positions == [pytest.approx(0.5)] + + final = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=11.5, dt=0.01)) + assert final is not None + assert final.joint_names == ["arm/joint2"] + assert trajectory_task.get_state() == TrajectoryState.COMPLETED + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=12.0, dt=0.01)) + is None + ) + + def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task): + first = JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + ], + ) + second = JointTrajectory( + joint_names=["arm/joint3"], + points=[ + TrajectoryPoint(positions=[2.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[3.0], velocities=[0.0], time_from_start=1.0), + ], + ) + assert trajectory_task.execute(first) is True + assert trajectory_task.execute(second) is True + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.0, dt=0.01)) + output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.5, dt=0.01)) + assert output is not None + assert output.joint_names == ["arm/joint3"] + assert trajectory_task.cancel() is True + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=2.0, dt=0.01)) + is None + ) + assert trajectory_task.reset() is True + assert trajectory_task.claim().joints == frozenset( + {"arm/joint1", "arm/joint2", "arm/joint3"} + ) + def test_compute_during_trajectory(self, trajectory_task, simple_trajectory, coordinator_state): t_start = time.perf_counter() trajectory_task.execute(simple_trajectory) diff --git a/dimos/manipulation/_test_manipulation_helpers.py b/dimos/manipulation/_test_manipulation_helpers.py new file mode 100644 index 0000000000..6df020d7bc --- /dev/null +++ b/dimos/manipulation/_test_manipulation_helpers.py @@ -0,0 +1,45 @@ +# 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. + +"""Shared lightweight test harnesses for manipulation module tests.""" + +import threading +from unittest.mock import MagicMock + +from dimos.manipulation.manipulation_module import ( + ManipulationModule, + ManipulationState, +) + + +class ManipulationModuleHarness(ManipulationModule): + """Manipulation module initialized only with state needed by unit tests.""" + + def __init__(self) -> None: + self._state = ManipulationState.IDLE + self._lock = threading.Lock() + self._error_message = "" + self._planning_epoch = 0 + self._robots = {} + self._last_plan = None + self._world_monitor = None + self._planner = None + self._kinematics = None + self._coordinator_client = None + self.config = MagicMock(planning_timeout=10.0) + + +def make_module() -> ManipulationModule: + """Create a lightweight ManipulationModule harness for behavior tests.""" + return ManipulationModuleHarness() diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index a0fed03f7b..25418b309c 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -26,11 +26,11 @@ from collections.abc import Mapping, Sequence from enum import Enum +import math import threading import time from typing import TYPE_CHECKING, Any, TypeAlias -import numpy as np from pydantic import Field from dimos.agents.annotation import skill @@ -51,7 +51,6 @@ filter_joint_state_to_selected_joints, joint_target_to_global_names, planning_group_id_from_selector, - project_global_joint_path_to_robot, ) from dimos.manipulation.planning.kinematics.config import ( ManipulationKinematicsConfig, @@ -63,7 +62,6 @@ from dimos.manipulation.planning.spec.models import ( GeneratedPlan, IKResult, - JointPath, Obstacle, PlanningGroupID, RobotName, @@ -79,6 +77,7 @@ NoManipulationVisualizationConfig, ) from dimos.manipulation.visualization.factory import create_manipulation_visualization +from dimos.manipulation.visualization.operator import ManipulationOperator from dimos.manipulation.visualization.types import TargetEvaluation from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -86,6 +85,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -256,7 +256,8 @@ def _initialize_planning(self) -> None: self._world_monitor.start_state_monitor(robot_id) self._world_monitor.set_visualization(visualization) - self._world_monitor.sync_visualization_scene() + operator = ManipulationOperator(self, self._world_monitor) + self._world_monitor.initialize_visualization(operator=operator) if self._world_monitor.visualization is not None: self._world_monitor.start_visualization_thread(rate_hz=10.0) @@ -342,7 +343,7 @@ def _on_joint_state(self, msg: JointState) -> None: # Route to specific monitor self._world_monitor.on_joint_state(sub_msg, robot_id=robot_id) - # Capture per-robot init joints on first receipt + # Capture per-robot init joints on first update if robot_name not in self._init_joints: self._init_joints[robot_name] = sub_msg logger.info( @@ -406,14 +407,17 @@ def get_error(self) -> str: @rpc def cancel(self) -> bool: """Cancel current motion or invalidate an in-progress plan.""" - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - logger.info("Planning cancelled") - return True - if self._state != ManipulationState.EXECUTING: - return False - self._state = ManipulationState.IDLE + with self._lock: + if self._state == ManipulationState.PLANNING: + self._planning_epoch += 1 + self._state = ManipulationState.IDLE + elif self._state == ManipulationState.EXECUTING: + self._state = ManipulationState.IDLE + else: + return False + plan = self._last_plan + if plan is not None: + self._dismiss_preview(plan.group_ids) logger.info("Motion cancelled") return True @@ -425,15 +429,16 @@ def reset(self) -> SkillResult[ManipulationSkillError]: Use this after an error or fault to allow new commands. Cannot reset while a motion is executing — cancel first. """ - if self._state == ManipulationState.EXECUTING: - return SkillResult.fail( - "INVALID_STATE", - "Cannot reset while executing — cancel the motion first", - ) - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - self._error_message = "" + with self._lock: + if self._state == ManipulationState.EXECUTING: + return SkillResult.fail( + "INVALID_STATE", + "Cannot reset while executing — cancel the motion first", + ) + if self._state == ManipulationState.PLANNING: + self._planning_epoch += 1 + self._state = ManipulationState.IDLE + self._error_message = "" return SkillResult.ok("Reset to IDLE — ready for new commands") @rpc @@ -496,6 +501,7 @@ def _begin_planning( logger.warning(f"Cannot plan: state is {self._state.name}") return None self._planning_epoch += 1 + self._last_plan = None self._state = ManipulationState.PLANNING return robot[0], robot[1] @@ -525,73 +531,162 @@ def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> Plan ) return group_id - def _project_plan_to_robot_paths( + def _split_plan_trajectory_by_robot( self, plan: GeneratedPlan - ) -> dict[RobotName, JointPath] | None: - """Project a selected-global-joint plan into full local paths per affected robot. - - TODO: need proper trajectory parametrization - """ - if self._world_monitor is None: - return None - paths: dict[RobotName, JointPath] = {} - try: - affected = self._world_monitor.planning_groups.select(plan.group_ids).robot_names - except Exception as exc: - logger.error("Failed to resolve generated plan: %s", exc) - return None - - for robot_name in affected: + ) -> dict[RobotName, JointTrajectory] | None: + """Split a stored global selected-joint trajectory into robot-local subsets.""" + trajectories: dict[RobotName, JointTrajectory] = {} + robot_indices: dict[RobotName, list[tuple[int, str]]] = {} + for index, global_name in enumerate(plan.trajectory.joint_names): + if "/" not in global_name: + logger.error("Stored trajectory joint '%s' is not globally named", global_name) + return None + robot_name, local_name = global_name.split("/", 1) robot = self._get_robot(robot_name) if robot is None: return None - resolved_name, robot_id, config, _ = robot - current = self._world_monitor.get_current_joint_state(robot_id) - try: - paths[resolved_name] = project_global_joint_path_to_robot( - plan.path, - robot_name=resolved_name, - local_joint_names=config.joint_names, - current_joint_state=current, - ) - except ValueError as exc: - logger.error("Cannot project plan for '%s': %s", resolved_name, exc) + _, _, config, _ = robot + if local_name not in config.joint_names: + logger.error("Stored trajectory joint '%s' is not configured", global_name) return None - return paths + robot_indices.setdefault(robot_name, []).append((index, local_name)) + + for robot_name, indexed_names in robot_indices.items(): + indices = [index for index, _ in indexed_names] + local_names = [name for _, name in indexed_names] + points = [ + TrajectoryPoint( + time_from_start=point.time_from_start, + positions=[point.positions[index] for index in indices], + velocities=[point.velocities[index] for index in indices], + ) + for point in plan.trajectory.points + ] + trajectories[robot_name] = JointTrajectory( + joint_names=local_names, + points=points, + timestamp=plan.trajectory.timestamp, + ) + logger.info( + "Stored trajectory for %s: %d selected joints, %.3fs", + robot_name, + len(local_names), + trajectories[robot_name].duration, + ) + return trajectories - def _trajectory_from_robot_path( - self, robot_name: RobotName, path: JointPath - ) -> JointTrajectory | None: - """Generate a local trajectory for one projected robot path.""" + @staticmethod + def _assert_finite_sequence(values: Sequence[float], label: str) -> None: + for value in values: + if not math.isfinite(value): + raise ValueError(f"{label} contains non-finite value") + + def _limits_for_global_joints( + self, joint_names: Sequence[str] + ) -> tuple[list[float], list[float]]: + velocities: list[float] = [] + accelerations: list[float] = [] + for global_name in joint_names: + if "/" not in global_name: + raise ValueError(f"Joint '{global_name}' is not globally named") + robot_name, local_name = global_name.split("/", 1) + robot = self._get_robot(robot_name) + if robot is None: + raise ValueError(f"Unknown robot for joint '{global_name}'") + _, _, config, _ = robot + if local_name not in config.joint_names: + raise ValueError(f"Unknown local joint '{global_name}'") + velocity = float(config.max_velocity) + acceleration = float(config.max_acceleration) + if not math.isfinite(velocity) or velocity <= 0.0: + raise ValueError(f"Invalid velocity limit for '{global_name}'") + if not math.isfinite(acceleration) or acceleration <= 0.0: + raise ValueError(f"Invalid acceleration limit for '{global_name}'") + velocities.append(velocity) + accelerations.append(acceleration) + return velocities, accelerations + + def _validate_selected_path( + self, path: Sequence[JointState], expected_names: Sequence[str] + ) -> list[list[float]]: if len(path) < 2: - logger.error("Plan projection for '%s' has fewer than two waypoints", robot_name) - return None - robot = self._get_robot(robot_name) - if robot is None: - return None - _, _, config, traj_gen = robot - local_trajectory = traj_gen.generate([list(state.position) for state in path]) - return JointTrajectory( - joint_names=list(config.joint_names), - points=local_trajectory.points, - timestamp=local_trajectory.timestamp, + raise ValueError("Planner returned fewer than two waypoints") + expected = list(expected_names) + waypoints: list[list[float]] = [] + for waypoint_index, state in enumerate(path): + if list(state.name) != expected: + raise ValueError( + f"Waypoint {waypoint_index} joint names do not match selected order" + ) + positions = list(state.position) + if len(positions) != len(expected): + raise ValueError(f"Waypoint {waypoint_index} position dimension mismatch") + self._assert_finite_sequence(positions, f"Waypoint {waypoint_index} positions") + waypoints.append(positions) + return waypoints + + def _validate_generated_trajectory( + self, + trajectory: JointTrajectory, + expected_names: Sequence[str], + waypoints: Sequence[Sequence[float]], + ) -> None: + expected = list(expected_names) + if list(trajectory.joint_names) != expected: + raise ValueError("Generated trajectory joint names do not match selected order") + if not trajectory.points: + raise ValueError("Generated trajectory has no points") + previous_time: float | None = None + for point_index, point in enumerate(trajectory.points): + if len(point.positions) != len(expected) or len(point.velocities) != len(expected): + raise ValueError(f"Generated point {point_index} dimension mismatch") + self._assert_finite_sequence( + point.positions, f"Generated point {point_index} positions" + ) + self._assert_finite_sequence( + point.velocities, f"Generated point {point_index} velocities" + ) + if not math.isfinite(point.time_from_start): + raise ValueError(f"Generated point {point_index} time is non-finite") + if point_index == 0 and point.time_from_start != 0.0: + raise ValueError("Generated trajectory must start at time 0") + if previous_time is not None and point.time_from_start <= previous_time: + raise ValueError("Generated trajectory times must be strictly increasing") + previous_time = point.time_from_start + non_noop = any(list(waypoint) != list(waypoints[0]) for waypoint in waypoints[1:]) + if non_noop and trajectory.duration <= 0.0: + raise ValueError("Generated trajectory duration must be positive") + waypoint_index = 0 + for point in trajectory.points: + if list(point.positions) == list(waypoints[waypoint_index]): + waypoint_index += 1 + if waypoint_index == len(waypoints): + break + if waypoint_index != len(waypoints): + raise ValueError("Generated trajectory does not contain ordered waypoint boundaries") + + def _materialize_generated_plan( + self, group_ids: tuple[PlanningGroupID, ...], result_path: Sequence[JointState] + ) -> tuple[list[JointState], JointTrajectory]: + assert self._world_monitor is not None + selection = self._world_monitor.planning_groups.select(group_ids) + expected_names = list(selection.joint_names) + path = [JointState(state) for state in result_path] + waypoints = self._validate_selected_path(path, expected_names) + velocities, accelerations = self._limits_for_global_joints(expected_names) + generator = JointTrajectoryGenerator( + num_joints=len(expected_names), + max_velocity=velocities, + max_acceleration=accelerations, ) - - def _trajectories_for_plan( - self, plan: GeneratedPlan - ) -> dict[RobotName, JointTrajectory] | None: - """Derive robot-local trajectories from a generated plan.""" - projected = self._project_plan_to_robot_paths(plan) - if projected is None: - return None - trajectories: dict[RobotName, JointTrajectory] = {} - for robot_name, path in projected.items(): - trajectory = self._trajectory_from_robot_path(robot_name, path) - if trajectory is None: - return None - trajectories[robot_name] = trajectory - logger.info("Trajectory for %s: %.3fs", robot_name, trajectory.duration) - return trajectories + generated = generator.generate(waypoints) + trajectory = JointTrajectory( + joint_names=expected_names, + points=generated.points, + timestamp=generated.timestamp, + ) + self._validate_generated_trajectory(trajectory, expected_names, waypoints) + return path, trajectory def _plan_selected_path( self, @@ -599,8 +694,8 @@ def _plan_selected_path( start: JointState, goal: JointState, planning_epoch: int, - ) -> bool: - """Plan over an explicit planning group selection and store the result.""" + ) -> GeneratedPlan | None: + """Plan over explicit planning groups and store the resulting plan.""" assert self._world_monitor and self._planner result = self._planner.plan_selected_joint_path( world=self._world_monitor.world, @@ -609,43 +704,71 @@ def _plan_selected_path( goal=goal, timeout=self.config.planning_timeout, ) - if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: - logger.info("Discarding cancelled planning result") - return False if not result.is_success(): - return self._fail(f"Planning failed: {result.status.name}") + _ = self._fail_planning_epoch(planning_epoch, f"Planning failed: {result.status.name}") + return None logger.info("Path: %d waypoints, groups=%s", len(result.path), group_ids) + try: + path, trajectory = self._materialize_generated_plan(group_ids, result.path) + except Exception as exc: + _ = self._fail_planning_epoch(planning_epoch, f"Failed to materialize plan: {exc}") + return None plan = GeneratedPlan( group_ids=group_ids, - path=result.path, + trajectory=trajectory, + path=path, status=result.status, planning_time=result.planning_time, path_length=result.path_length, iterations=result.iterations, message=result.message, ) - projected = self._project_plan_to_robot_paths(plan) - if projected is None or any(len(path) < 2 for path in projected.values()): - self._last_plan = None - return self._fail("Failed to project generated plan") - self._last_plan = plan - self._state = ManipulationState.COMPLETED - return True + with self._lock: + if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: + logger.info("Discarding cancelled planning result") + return None + self._last_plan = plan + self._state = ManipulationState.COMPLETED + return plan def _fail(self, msg: str) -> bool: """Set FAULT state with error message.""" logger.warning(msg) - self._state = ManipulationState.FAULT - self._error_message = msg + with self._lock: + self._state = ManipulationState.FAULT + self._error_message = msg return False - def _dismiss_preview(self, robot_id: WorldRobotID) -> None: + def _fail_planning_epoch(self, planning_epoch: int, msg: str) -> bool: + """Fault only the still-current planning operation.""" + with self._lock: + if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: + logger.info("Discarding cancelled planning result") + return False + logger.warning(msg) + self._last_plan = None + self._state = ManipulationState.FAULT + self._error_message = msg + return False + + def _dismiss_preview(self, group_ids: Sequence[PlanningGroupID]) -> None: """Hide the preview ghost if the world supports it.""" if self._world_monitor is None: return - self._world_monitor.hide_preview(robot_id) - self._world_monitor.publish_visualization() + try: + robot_names = self._world_monitor.planning_groups.select(tuple(group_ids)).robot_names + robot_ids = tuple( + robot_id + for robot_name in robot_names + if (robot_id := self.robot_id_for_name(robot_name)) is not None + ) + except (KeyError, ValueError): + robot_ids = () + if robot_ids: + self._world_monitor.cancel_preview_animation(robot_ids=robot_ids) + else: + self._world_monitor.cancel_preview_animation() def _solve_ik_for_pose( self, @@ -800,41 +923,7 @@ def plan_to_pose_targets( auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), ) -> bool: """Plan to one or more group pose targets with optional auxiliary groups.""" - if self._world_monitor is None or self._kinematics is None: - return False - if not pose_targets: - return self._fail("At least one pose target is required") - - stamped_targets = { - planning_group_id_from_selector(group): PoseStamped( - frame_id="world", - position=pose.position, - orientation=pose.orientation, - ) - for group, pose in pose_targets.items() - } - auxiliary_ids = tuple(planning_group_id_from_selector(group) for group in auxiliary_groups) - group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) - planning_epoch = self._begin_group_planning() - if planning_epoch is None: - return False - - try: - selection = self._world_monitor.planning_groups.select(group_ids) - current = self._world_monitor.current_global_joint_state() - start = filter_joint_state_to_selected_joints(current, selection.joint_names) - except Exception as exc: - return self._fail(f"Failed to resolve planning groups: {exc}") - - ik = self.inverse_kinematics( - pose_targets=stamped_targets, - auxiliary_group_ids=auxiliary_ids, - seed=start, - ) - if not ik.is_success() or ik.joint_state is None: - return self._fail(f"IK failed: {ik.status.name}") - logger.info(f"IK solved, error: {ik.position_error:.4f}m") - return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) + return self.generate_plan_to_pose_targets(pose_targets, auxiliary_groups) is not None @rpc def plan_to_joints(self, joints: JointState, robot_name: RobotName | None = None) -> bool: @@ -869,24 +958,34 @@ def plan_to_joint_targets( self, joint_targets: Mapping[PlanningGroupID | PlanningGroup, JointState] ) -> bool: """Plan to joint targets keyed by planning group.""" + return self.generate_plan_to_joint_targets(joint_targets) is not None + + def generate_plan_to_joint_targets( + self, joint_targets: Mapping[PlanningGroupID | PlanningGroup, JointState] + ) -> GeneratedPlan | None: + """Plan to joint targets and return the exact stored GeneratedPlan.""" if self._world_monitor is None or self._planner is None: - return False + return None if not joint_targets: - return self._fail("At least one joint target is required") + _ = self._fail("At least one joint target is required") + return None group_ids = tuple( dict.fromkeys(planning_group_id_from_selector(group) for group in joint_targets) ) planning_epoch = self._begin_group_planning() if planning_epoch is None: - return False + return None try: selection = self._world_monitor.planning_groups.select(group_ids) current = self._world_monitor.current_global_joint_state() start = filter_joint_state_to_selected_joints(current, selection.joint_names) except Exception as exc: - return self._fail(f"Failed to resolve planning groups: {exc}") + _ = self._fail_planning_epoch( + planning_epoch, f"Failed to resolve planning groups: {exc}" + ) + return None goal_names: list[str] = [] goal_positions: list[float] = [] @@ -897,48 +996,59 @@ def plan_to_joint_targets( target_global = joint_target_to_global_names(target_group, target) except (KeyError, ValueError) as exc: logger.error(str(exc)) - return self._fail(f"Invalid joint target for '{group_id}'") + _ = self._fail_planning_epoch( + planning_epoch, f"Invalid joint target for '{group_id}'" + ) + return None goal_names.extend(target_global.name) goal_positions.extend(target_global.position) goal = JointState(name=goal_names, position=goal_positions) return self._plan_selected_path(group_ids, start, goal, planning_epoch) - def _preview_projected_path( + def generate_plan_to_pose_targets( self, - robot_id: WorldRobotID, - planned_path: JointPath, - trajectory: JointTrajectory, - duration: float | None, - target_fps: float, - ) -> bool: - """Animate a projected local path with optional trajectory interpolation.""" - if self._world_monitor is None or len(planned_path) == 0: - return False - - animation_duration = duration if duration is not None else trajectory.duration - interpolated = list(planned_path) - if target_fps > 0 and animation_duration > 0: - times = np.array( - [point.time_from_start for point in trajectory.points], dtype=np.float64 + pose_targets: Mapping[PlanningGroupID | PlanningGroup, Pose], + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> GeneratedPlan | None: + """Plan to pose targets and return the exact stored GeneratedPlan.""" + if self._world_monitor is None or self._kinematics is None: + return None + if not pose_targets: + _ = self._fail("At least one pose target is required") + return None + stamped_targets = { + planning_group_id_from_selector(group): PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, ) - positions = np.array([point.positions for point in trajectory.points], dtype=np.float64) - if len(times) > 1 and positions.ndim == 2 and times[-1] > times[0]: - frame_count = int(np.ceil(animation_duration * target_fps)) + 1 - sample_times = np.linspace(times[0], times[-1], frame_count) - joint_names = trajectory.joint_names or planned_path[0].name - sampled_positions = np.column_stack( - [ - np.interp(sample_times, times, positions[:, joint]) - for joint in range(positions.shape[1]) - ] - ) - interpolated = [ - JointState(name=joint_names, position=position.tolist()) - for position in sampled_positions - ] - self._world_monitor.animate_path(robot_id, interpolated, animation_duration) - return True + for group, pose in pose_targets.items() + } + auxiliary_ids = tuple(planning_group_id_from_selector(group) for group in auxiliary_groups) + group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) + planning_epoch = self._begin_group_planning() + if planning_epoch is None: + return None + try: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + except Exception as exc: + _ = self._fail_planning_epoch( + planning_epoch, f"Failed to resolve planning groups: {exc}" + ) + return None + ik = self.inverse_kinematics( + pose_targets=stamped_targets, + auxiliary_group_ids=auxiliary_ids, + seed=start, + ) + if not ik.is_success() or ik.joint_state is None: + _ = self._fail_planning_epoch(planning_epoch, f"IK failed: {ik.status.name}") + return None + logger.info(f"IK solved, error: {ik.position_error:.4f}m") + return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) @rpc def preview_path( @@ -950,9 +1060,9 @@ def preview_path( """Compatibility wrapper for preview_plan(). Args: - duration: Total animation duration in seconds. Uses trajectory duration if None. - robot_name: Robot to preview (required if multiple robots configured) - target_fps: Nominal preview update rate. Set <= 0 to use planned waypoints directly. + duration: Total animation duration in seconds. Defaults to one second. + robot_name: Compatibility affected-robot validation; does not filter the preview. + target_fps: Deprecated compatibility argument; shared-clock previews use plan waypoints. """ return self.preview_plan(None, duration, robot_name, target_fps) @@ -964,14 +1074,14 @@ def preview_plan( robot_name: RobotName | None = None, target_fps: float = 30.0, ) -> bool: - """Preview a generated plan in the visualizer.""" + """Preview a complete generated plan in the visualizer.""" plan = plan or self._last_plan if plan is None or not plan.path: logger.warning("No generated plan to preview") return False try: assert self._world_monitor is not None - affected = list(self._world_monitor.planning_groups.select(plan.group_ids).robot_names) + affected = self._world_monitor.planning_groups.select(plan.group_ids).robot_names except Exception as exc: logger.error("Generated plan cannot be resolved: %s", exc) return False @@ -979,20 +1089,10 @@ def preview_plan( if robot_name not in affected: logger.error("Generated plan does not affect robot '%s'", robot_name) return False - affected = [robot_name] - projected = self._project_plan_to_robot_paths(plan) - if projected is None: + if self._world_monitor is None: return False - for name in affected: - robot = self._get_robot(name) - if robot is None or (planned_path := projected.get(name)) is None: - return False - _, robot_id, _, _ = robot - trajectory = self._trajectory_from_robot_path(name, planned_path) - if trajectory is None or not self._preview_projected_path( - robot_id, planned_path, trajectory, duration, target_fps - ): - return False + _ = target_fps + self._world_monitor.animate_trajectory(plan.trajectory, duration) return True @rpc @@ -1022,9 +1122,16 @@ def clear_planned_path(self) -> bool: Returns: True if cleared """ - if self._world_monitor is None: - return False - self._last_plan = None + with self._lock: + plan = self._last_plan + self._last_plan = None + if self._state == ManipulationState.PLANNING: + self._planning_epoch += 1 + self._state = ManipulationState.IDLE + if plan is not None: + # Preserve the group selection until the public visualization + # transaction has invalidated and hidden its preview. + self._dismiss_preview(plan.group_ids) return True @rpc @@ -1326,57 +1433,159 @@ def execute_plan( self, plan: GeneratedPlan | None = None, robot_name: RobotName | None = None ) -> bool: """Execute a generated planning-group plan through affected trajectory tasks.""" - if self._world_monitor is None: - logger.error("Planning not initialized") + execution = self._begin_execution(plan) + if execution is None: return False - plan = plan or self._last_plan - if plan is None or not plan.path: - logger.warning("No generated plan") - return False - if (client := self._get_coordinator_client()) is None: - logger.error("No coordinator client") + previous_state, target_plan = execution + try: + freshness_error = self._stored_plan_freshness_error(target_plan) + if freshness_error is not None: + self._restore_execution_gate(previous_state, freshness_error) + return False + prepared = self._prepare_execution(target_plan, robot_name) + if isinstance(prepared, str): + self._restore_execution_gate(previous_state, prepared) + return False + except Exception as exc: + self._restore_execution_gate(previous_state, f"Failed to prepare execution: {exc}") return False + return self._dispatch_prepared_plan(prepared) + + def _begin_execution( + self, plan: GeneratedPlan | None + ) -> tuple[ManipulationState, GeneratedPlan] | None: + """Enter EXECUTING if the module is idle enough to dispatch.""" + with self._lock: + target_plan = plan or self._last_plan + if target_plan is None or not target_plan.path or not target_plan.trajectory.points: + logger.warning("Stored plan is invalid or not executable") + self._error_message = "Stored plan is invalid or not executable" + return None + if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): + logger.warning("Manipulation state is not executable") + return None + previous_state = self._state + self._state = ManipulationState.EXECUTING + return previous_state, target_plan + + def _restore_execution_gate(self, previous_state: ManipulationState, msg: str) -> None: + logger.warning(msg) + with self._lock: + self._state = previous_state + self._error_message = msg + + def _stored_plan_freshness_error( + self, plan: GeneratedPlan, tolerance: float = 1e-6 + ) -> str | None: + """Return why current selected joints no longer match the plan start.""" + if self._world_monitor is None: + return "Planning not initialized" + if not plan.trajectory.points: + return "Stored trajectory has no points" + first = plan.trajectory.points[0] + names = list(plan.trajectory.joint_names) + if len(names) != len(first.positions): + return "Stored trajectory start has malformed joint positions" + try: + current = self._world_monitor.current_global_joint_state() + except (AttributeError, TypeError, ValueError): + return "Current planned joints are malformed" + if not isinstance(current, JointState) or len(current.name) != len(current.position): + return "Current planned joints are malformed" + values: dict[str, float] = {} + ordered_subset: list[tuple[str, float]] = [] + planned = set(names) + for raw_name, raw_position in zip(current.name, current.position, strict=True): + name = str(raw_name) + if name in values: + return "Current planned joints are malformed" + try: + value = float(raw_position) + except (TypeError, ValueError): + return f"Current planned joint '{name}' is malformed" + if not math.isfinite(value): + return f"Current planned joint '{name}' is malformed" + values[name] = value + if name in planned: + ordered_subset.append((name, value)) + if len(ordered_subset) != len(names): + return "Current planned joints are malformed" + for (actual_name, actual), name, expected in zip( + ordered_subset, names, first.positions, strict=True + ): + if actual_name != name: + return "Current planned joints are not in stored plan order" + if not math.isfinite(actual) or not math.isfinite(float(expected)): + return f"Current planned joint '{name}' is malformed" + if abs(actual - float(expected)) > tolerance: + return f"Current planned joint '{name}' no longer matches the stored plan start" + return None + + def _prepare_execution( + self, + plan: GeneratedPlan, + robot_name: RobotName | None = None, + ) -> tuple[tuple[RobotName, RobotModelConfig, JointTrajectory], ...] | str: + """Resolve execution inputs before dispatching any coordinator task.""" + if self._world_monitor is None: + return "Planning not initialized" + if not plan.path: + return "No generated plan" try: affected = list(self._world_monitor.planning_groups.select(plan.group_ids).robot_names) except Exception as exc: - return self._fail(f"Failed to resolve generated plan: {exc}") + return f"Failed to resolve generated plan: {exc}" if robot_name is not None: if robot_name not in affected: - logger.warning("No planned path for '%s'", robot_name) - return False + return f"No planned path for '{robot_name}'" affected = [robot_name] - trajectories = self._trajectories_for_plan(plan) + trajectories = self._split_plan_trajectory_by_robot(plan) if trajectories is None: - return self._fail("Failed to project generated plan") - - self._state = ManipulationState.EXECUTING - for robot_name in affected: - robot = self._get_robot(robot_name) + return "Failed to split stored trajectory" + if self._get_coordinator_client() is None: + return "No coordinator client" + prepared: list[tuple[RobotName, RobotModelConfig, JointTrajectory]] = [] + for name in affected: + robot = self._get_robot(name) if robot is None: - return self._fail(f"Robot '{robot_name}' not found") + return f"Robot '{name}' not found" _, _, config, _ = robot if not config.coordinator_task_name: - logger.error(f"No coordinator_task_name for '{robot_name}'") - return self._fail(f"No coordinator_task_name for '{robot_name}'") - traj = trajectories.get(robot_name) + return f"No coordinator_task_name for '{name}'" + traj = trajectories.get(name) if traj is None: - logger.warning("No planned trajectory for '%s'", robot_name) - return self._fail(f"No planned trajectory for '{robot_name}'") - translated = self._translate_trajectory_to_coordinator(traj, config) - logger.info( - "Executing: task='%s', %d pts, %.2fs", - config.coordinator_task_name, - len(translated.points), - translated.duration, - ) - result = client.task_invoke( - config.coordinator_task_name, "execute", {"trajectory": translated} - ) + return f"No planned trajectory for '{name}'" + prepared.append((name, config, traj)) + return tuple(prepared) + + def _dispatch_prepared_plan( + self, prepared: Sequence[tuple[RobotName, RobotModelConfig, JointTrajectory]] + ) -> bool: + """Dispatch already-validated robot trajectories; failures here fault.""" + if (client := self._get_coordinator_client()) is None: + logger.error("No coordinator client") + return False + + for _robot_name, config, traj in prepared: + try: + translated = self._translate_trajectory_to_coordinator(traj, config) + logger.info( + "Executing: task='%s', %d pts, %.2fs", + config.coordinator_task_name, + len(translated.points), + translated.duration, + ) + result = client.task_invoke( + config.coordinator_task_name, "execute", {"trajectory": translated} + ) + except Exception as exc: + return self._fail(f"Failed to dispatch trajectory: {exc}") if not result: return self._fail("Coordinator rejected trajectory") logger.info("Trajectory accepted") - self._state = ManipulationState.COMPLETED + with self._lock: + self._state = ManipulationState.COMPLETED return True @rpc @@ -1568,13 +1777,9 @@ def _wait_for_trajectory_completion( robot = self._get_robot(robot_name) if robot is None: return True - rname, _, config, _ = robot + _, _, config, _ = robot client = self._get_coordinator_client() - trajectory: JointTrajectory | None = None - if last_plan is not None and last_plan.path: - projected = self._project_plan_to_robot_paths(last_plan) - if projected is not None and (path := projected.get(rname)) is not None: - trajectory = self._trajectory_from_robot_path(rname, path) + trajectory = last_plan.trajectory if last_plan is not None else None if client is None or not config.coordinator_task_name: # No coordinator — wait for trajectory duration as fallback diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index a78957858c..a9781a6c8a 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -24,11 +24,17 @@ from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.monitor import world_monitor as world_monitor_module from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) +from dimos.manipulation.planning.spec.protocols import VisualizationSpec 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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class _VectorLike(list[float]): @@ -151,19 +157,16 @@ def get_group_jacobian(self, ctx, group_id): def get_visualization_url(self): return None - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: return None - def publish_visualization(self, ctx=None) -> None: + def update_state(self, frame: VisualizationStateFrame) -> None: return None - def show_preview(self, robot_id) -> None: + def animate_trajectory(self, trajectory, duration: float | None = None) -> None: return None - def hide_preview(self, robot_id) -> None: - return None - - def animate_path(self, robot_id, path, duration: float = 3.0) -> None: + def cancel_preview_animation(self) -> None: return None def close(self) -> None: @@ -177,20 +180,17 @@ def __init__(self) -> None: def get_visualization_url(self): return None - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - self.calls.append(("initialize_scene", scene)) - - def publish_visualization(self, ctx=None) -> None: - return None + def initialize(self, session: VisualizationSession) -> None: + self.calls.append(("initialize", session)) - def show_preview(self, robot_id) -> None: - self.calls.append(("show_preview", robot_id)) + def update_state(self, frame: VisualizationStateFrame) -> None: + self.calls.append(("update_state", frame)) - def hide_preview(self, robot_id) -> None: - self.calls.append(("hide_preview", robot_id)) + def animate_trajectory(self, trajectory, duration: float | None = None) -> None: + self.calls.append(("animate_trajectory", trajectory, duration)) - def animate_path(self, robot_id, path, duration: float = 3.0) -> None: - return None + def cancel_preview_animation(self) -> None: + self.calls.append(("cancel_preview_animation",)) def close(self) -> None: self.calls.append(("close", None)) @@ -251,10 +251,28 @@ def test_world_monitor_syncs_planning_scene_to_visualization() -> None: monitor.add_robot(_robot_config()) monitor.sync_visualization_scene() - assert fake_viz.calls[0][0] == "initialize_scene" - scene = fake_viz.calls[0][1] + assert fake_viz.calls[0][0] == "initialize" + session = fake_viz.calls[0][1] + scene = session.scene assert isinstance(scene, PlanningSceneInfo) assert scene.robots["robot-1"].name == "arm" + assert scene.planning_groups[0].id == "arm/manipulator" + + +def test_world_monitor_forwards_raw_trajectory_preview_protocol() -> None: + fake_viz = FakeViz() + monitor = world_monitor_module.WorldMonitor(world=FakeWorld(), visualization=fake_viz) # type: ignore[arg-type] + trajectory = JointTrajectory(joint_names=["arm/j1"], points=[]) + + assert isinstance(fake_viz, VisualizationSpec) + monitor.cancel_preview_animation() + monitor.animate_trajectory(trajectory, 2.0) + + assert fake_viz.calls == [ + ("cancel_preview_animation",), + ("cancel_preview_animation",), + ("animate_trajectory", trajectory, 2.0), + ] def test_create_planning_specs_wraps_existing_world(monkeypatch) -> None: diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 9b54332de8..ed4da16837 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections.abc import Sequence from contextlib import contextmanager import threading from typing import TYPE_CHECKING, Any @@ -29,10 +30,15 @@ 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.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -67,6 +73,9 @@ def __init__( self._world = world self._visualization = visualization self._lock = threading.RLock() + # Keep renderer mutations and periodic publishes ordered. Cancellation is + # deliberately issued outside this lock so it can interrupt an animation. + self._visualization_lock = threading.RLock() self._robot_joints: dict[WorldRobotID, list[str]] = {} self._robot_configs: dict[WorldRobotID, RobotModelConfig] = {} self._robot_ids_by_name: dict[RobotName, WorldRobotID] = {} @@ -101,14 +110,23 @@ def planning_groups(self) -> PlanningGroupRegistry: def planning_scene_info(self) -> PlanningSceneInfo: """Return a stable metadata snapshot of the initialized planning scene.""" with self._lock: - return PlanningSceneInfo(robots=dict(self._robot_configs)) + return PlanningSceneInfo( + robots=dict(self._robot_configs), + planning_groups=tuple(self._planning_groups.list()), + ) - def sync_visualization_scene(self) -> None: - """Synchronize startup scene metadata to the attached visualization.""" + def initialize_visualization(self, operator: object | None = None) -> None: + """Initialize attached visualization with immutable startup metadata.""" visualization = self._visualization if visualization is None: return - visualization.initialize_scene(self.planning_scene_info()) + visualization.initialize( + VisualizationSession(scene=self.planning_scene_info(), operator=operator) + ) + + def sync_visualization_scene(self) -> None: + """Compatibility wrapper for initializing visualization metadata.""" + self.initialize_visualization() def get_robot_ids(self) -> list[WorldRobotID]: """Get all robot IDs.""" @@ -214,7 +232,10 @@ def stop_all_monitors(self) -> None: logger.info("All monitors stopped") if self._visualization is not None: - self._visualization.close() + self._visualization.cancel_preview_animation() + # Wait for cancelled animation cleanup before releasing renderer resources. + with self._visualization_lock: + self._visualization.close() # Message Handlers @@ -522,30 +543,57 @@ def get_visualization_url(self) -> str | None: return str(url) if url else None return None - def publish_visualization(self) -> None: - """Force publish current state to visualization.""" - if self._visualization is not None: - self._visualization.publish_visualization() + def visualization_state_frame(self) -> VisualizationStateFrame: + """Build a pushed visualization state frame without freshness policy.""" + joint_states: dict[str, JointState] = {} + with self._lock: + robot_ids = list(self._robot_configs.keys()) + for robot_id in robot_ids: + state = self.get_current_joint_state(robot_id) + if state is not None: + joint_states[robot_id] = state + return VisualizationStateFrame(joint_states=joint_states) - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview representation for a robot if visualization is available.""" + def update_visualization_state(self) -> None: + """Push current state to visualization.""" if self._visualization is not None: - self._visualization.show_preview(robot_id) + with self._visualization_lock: + self._visualization.update_state(self.visualization_state_frame()) - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview representation for a robot if visualization is available.""" + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Cancel active visualization preview animation.""" if self._visualization is not None: - self._visualization.hide_preview(robot_id) + if robot_ids is None: + self._visualization.cancel_preview_animation() + else: + self._visualization.cancel_preview_animation(robot_ids) - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: - """Animate a path if visualization is available.""" + """Animate a raw generated-plan trajectory if visualization is available.""" if self._visualization is not None: - self._visualization.animate_path(robot_id, path, duration) + robot_ids = self.robot_ids_for_global_joints(trajectory.joint_names) + if robot_ids: + self._visualization.cancel_preview_animation(robot_ids) + else: + self._visualization.cancel_preview_animation() + with self._visualization_lock: + self._visualization.animate_trajectory(trajectory, duration) + + def robot_ids_for_global_joints(self, joint_names: Sequence[str]) -> tuple[WorldRobotID, ...]: + """Return visualization robot IDs affected by globally named trajectory joints.""" + robot_ids: list[WorldRobotID] = [] + with self._lock: + by_name = {config.name: robot_id for robot_id, config in self._robot_configs.items()} + for joint_name in joint_names: + if "/" not in joint_name: + continue + robot_name, _ = joint_name.split("/", 1) + robot_id = by_name.get(robot_name) + if robot_id is not None and robot_id not in robot_ids: + robot_ids.append(robot_id) + return tuple(robot_ids) def start_visualization_thread(self, rate_hz: float = 10.0) -> None: """Start background thread for visualization updates at given rate.""" @@ -586,7 +634,7 @@ def _visualization_loop(self) -> None: period = 1.0 / self._viz_rate_hz while not self._viz_stop_event.is_set(): try: - self.publish_visualization() + self.update_visualization_state() except Exception as e: logger.debug(f"Visualization publish failed: {e}") time.sleep(period) diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index 2dfec92e2d..b83dbbc2e6 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -30,9 +30,11 @@ import numpy as np from numpy.typing import NDArray + from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState + from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory RobotName: TypeAlias = str @@ -65,6 +67,25 @@ class PlanningSceneInfo: robots: Mapping[WorldRobotID, RobotModelConfig] """Robot model configurations keyed by world robot ID.""" + planning_groups: tuple[PlanningGroup, ...] = () + """Resolved immutable planning groups for the initialized scene.""" + + +@dataclass(frozen=True) +class VisualizationSession: + """One-shot immutable visualization initialization payload.""" + + scene: PlanningSceneInfo + operator: object | None = None + """Optional concrete ManipulationOperator; typed as object to avoid low-level cycles.""" + + +@dataclass(frozen=True) +class VisualizationStateFrame: + """Pushed current joint states for visualization backends.""" + + joint_states: Mapping[WorldRobotID, JointState] + Jacobian: TypeAlias = "NDArray[np.float64]" """6 x n Jacobian matrix (rows: [vx, vy, vz, wx, wy, wz])""" @@ -152,9 +173,10 @@ def is_success(self) -> bool: @dataclass class GeneratedPlan: - """Canonical selected-planning-group path exposed by ManipulationModule.""" + """Canonical selected-planning-group plan exposed by ManipulationModule.""" group_ids: tuple[PlanningGroupID, ...] + trajectory: JointTrajectory path: list[JointState] = field(default_factory=list) status: PlanningStatus = PlanningStatus.NO_SOLUTION planning_time: float = 0.0 diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 9ea84ab646..f8762521d7 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -33,15 +33,16 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( IKResult, - JointPath, Obstacle, PlanningGroupID, PlanningResult, - PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, WorldRobotID, ) from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState + from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory @runtime_checkable @@ -193,28 +194,26 @@ class VisualizationSpec(Protocol): visualization affordances. """ - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - """Receive stable planning-scene metadata after world startup.""" + def initialize(self, session: VisualizationSession) -> None: + """Receive one-shot visualization session metadata after world startup.""" ... def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" ... - def publish_visualization(self, ctx: Any | None = None) -> None: - """Publish current state to visualization.""" + def update_state(self, frame: VisualizationStateFrame) -> None: + """Receive current joint states keyed by initialized world robot ID.""" ... - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview representation for a robot.""" + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: + """Animate a raw globally named trajectory.""" ... - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview representation for a robot.""" - ... - - def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0) -> None: - """Animate a path in visualization.""" + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Cancel an active preview animation without waiting for its renderer to finish.""" ... def close(self) -> None: diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 2b4ad6f971..09bb30dc80 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field @@ -34,13 +35,7 @@ from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType -from dimos.manipulation.planning.spec.models import ( - JointPath, - Obstacle, - PlanningGroupID, - PlanningSceneInfo, - WorldRobotID, -) +from dimos.manipulation.planning.spec.models import Obstacle, PlanningGroupID, WorldRobotID from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.utils.logging_config import setup_logger @@ -53,6 +48,13 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + +if TYPE_CHECKING: + from dimos.manipulation.planning.spec.models import ( + VisualizationSession, + VisualizationStateFrame, + ) try: from pydrake.geometry import ( @@ -202,6 +204,8 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: self._plant_context: Context | None = None self._scene_graph_context: Context | None = None self._finalized = False + self._preview_animation_generation = 0 + self._preview_animation_generations: dict[WorldRobotID, int] = {} # Obstacle source for dynamic obstacles self._obstacle_source_id: Any = None @@ -786,10 +790,10 @@ def finalize(self) -> None: # Initial visualization publish (routed to Meshcat thread) if self._meshcat_visualizer is not None: - self.publish_visualization() + self._publish_visualization() # Hide all preview robots initially for robot_id in self._robots: - self.hide_preview(robot_id) + self._set_preview_visibility(robot_id, False) @property def is_finalized(self) -> bool: @@ -1151,8 +1155,9 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray # Visualization - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: """Embedded Meshcat observes the Drake world directly; no extra sync needed.""" + _ = session return None def get_visualization_url(self) -> str | None: @@ -1161,7 +1166,7 @@ def get_visualization_url(self) -> str | None: return self._meshcat.web_url() return None - def publish_visualization(self, ctx: Context | None = None) -> None: + def _publish_visualization(self, ctx: Context | None = None) -> None: """Publish current state to visualization.""" if self._meshcat_visualizer is None or self._meshcat is None: return @@ -1171,6 +1176,11 @@ def publish_visualization(self, ctx: Context | None = None) -> None: viz_ctx = self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) self._meshcat.forced_publish(self._meshcat_visualizer, viz_ctx) + def update_state(self, frame: VisualizationStateFrame) -> None: + """Receive pushed state frame; embedded Meshcat uses Drake live context.""" + _ = frame + self._publish_visualization() + def _set_preview_positions( self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] ) -> None: @@ -1184,57 +1194,127 @@ def _set_preview_positions( full_positions[idx] = positions[i] self._plant.SetPositions(plant_ctx, full_positions) - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview (yellow ghost) robot in Meshcat.""" + def _set_preview_visibility(self, robot_id: WorldRobotID, visible: bool) -> None: + """Set one preview robot's Meshcat visibility.""" if self._meshcat is None: return robot_data = self._robots.get(robot_id) if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", True) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", visible) - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview (yellow ghost) robot in Meshcat.""" - if self._meshcat is None: - return - robot_data = self._robots.get(robot_id) - if robot_data is None or robot_data.preview_model_instance is None: - return - model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", False) - - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Invalidate active preview frames and hide preview ghosts immediately.""" + with self._lock: + self._preview_animation_generation += 1 + affected = set(robot_ids) if robot_ids is not None else set(self._robots) + for robot_id in affected: + self._preview_animation_generations[robot_id] = ( + self._preview_animation_generations.get(robot_id, 0) + 1 + ) + if robot_id not in self._robots: + continue + self._set_preview_visibility(robot_id, False) + + def _robot_trajectory_indices( + self, trajectory: JointTrajectory + ) -> dict[WorldRobotID, list[tuple[int, str]]]: + robot_ids_by_name = { + robot.config.name: robot_id for robot_id, robot in self._robots.items() + } + indices: dict[WorldRobotID, list[tuple[int, str]]] = {} + for index, global_name in enumerate(trajectory.joint_names): + if "/" not in global_name: + raise ValueError(f"trajectory joint '{global_name}' is not globally named") + robot_name, local_name = global_name.split("/", 1) + robot_id = robot_ids_by_name.get(robot_name) + if robot_id is None: + raise ValueError(f"trajectory references unknown robot '{robot_name}'") + if local_name not in self._robots[robot_id].config.joint_names: + raise ValueError(f"trajectory references unknown joint '{global_name}'") + indices.setdefault(robot_id, []).append((index, local_name)) + return indices + + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: - """Animate a path using the preview (yellow ghost) robot. - - The preview stays visible after animation completes. - """ - if self._meshcat is None or len(path) < 2: - return - - robot_data = self._robots.get(robot_id) - if robot_data is None or robot_data.preview_model_instance is None: + """Render raw globally named trajectory on its stored shared clock.""" + if self._meshcat is None or len(trajectory.points) < 2: return import time - self.show_preview(robot_id) - dt = duration / (len(path) - 1) - for joint_state in path: - positions = np.array(joint_state.position, dtype=np.float64) + robot_indices = self._robot_trajectory_indices(trajectory) + robot_ids = list(robot_indices) + playback_scale = 1.0 + if duration is not None: + if duration <= 0.0 or trajectory.duration <= 0.0: + raise ValueError("preview duration must be positive") + playback_scale = duration / trajectory.duration + baselines: dict[WorldRobotID, NDArray[np.float64]] = {} + joint_positions_by_name: dict[WorldRobotID, dict[str, int]] = {} + with self._lock: + assert self._plant_context is not None + assert self._live_context is not None + self._preview_animation_generation += 1 + animation_generations: dict[WorldRobotID, int] = {} + for robot_id in robot_ids: + self._preview_animation_generations[robot_id] = ( + self._preview_animation_generations.get(robot_id, 0) + 1 + ) + animation_generations[robot_id] = self._preview_animation_generations[robot_id] + robot_data = self._robots[robot_id] + self._set_preview_visibility(robot_id, True) + baselines[robot_id] = np.array( + self.get_joint_state(self._live_context, robot_id).position, + dtype=np.float64, + ) + joint_positions_by_name[robot_id] = dict( + zip( + robot_data.config.joint_names, + range(len(robot_data.config.joint_names)), + strict=True, + ) + ) + + try: + previous_time = trajectory.points[0].time_from_start + for frame_index, point in enumerate(trajectory.points): + with self._lock: + active_robot_ids = [ + robot_id + for robot_id in robot_ids + if self._preview_animation_generations.get(robot_id) + == animation_generations[robot_id] + ] + if not active_robot_ids: + return + assert self._plant_context is not None + for robot_id in active_robot_ids: + indexed_names = robot_indices[robot_id] + positions = baselines[robot_id].copy() + local_index = joint_positions_by_name[robot_id] + for trajectory_index, local_name in indexed_names: + positions[local_index[local_name]] = point.positions[trajectory_index] + self._set_preview_positions(self._plant_context, robot_id, positions) + self._publish_visualization() + if frame_index < len(trajectory.points) - 1: + next_time = trajectory.points[frame_index + 1].time_from_start + time.sleep((next_time - previous_time) * playback_scale) + previous_time = next_time + finally: with self._lock: - assert self._plant_context is not None - self._set_preview_positions(self._plant_context, robot_id, positions) - self.publish_visualization() - time.sleep(dt) + for robot_id in robot_ids: + if ( + self._preview_animation_generations.get(robot_id) + == animation_generations[robot_id] + ): + self._set_preview_visibility(robot_id, False) def close(self) -> None: """Shut down the viz thread.""" + self.cancel_preview_animation() if self._meshcat is not None: self._meshcat.close() diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index de4c31ee95..7832626fa5 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -24,6 +24,8 @@ from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint requires_drake = pytest.mark.skipif( not DRAKE_AVAILABLE, @@ -31,6 +33,16 @@ ) +def _trajectory(names: list[str], first: list[float], second: list[float]) -> JointTrajectory: + return JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint(time_from_start=0.0, positions=first, velocities=[0.0] * len(names)), + TrajectoryPoint(time_from_start=2.0, positions=second, velocities=[0.0] * len(names)), + ], + ) + + def _write_urdf(path: Path) -> None: path.write_text( """ @@ -258,3 +270,180 @@ def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Pa with pytest.raises(ValueError, match="non-controllable"): world.get_group_jacobian(world.get_live_context(), "arm/arm") + + +@requires_drake +def test_drake_animate_trajectory_projects_all_robots_on_shared_ticks( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + left_config = _config(urdf, [_arm_group("joint1")]).model_copy(update={"name": "left"}) + right_config = _config(urdf, [_arm_group("joint2")]).model_copy(update={"name": "right"}) + left_id = world.add_robot(left_config) + right_id = world.add_robot(right_config) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + ctx = world.get_live_context() + world.set_joint_state(ctx, left_id, JointState(name=["joint1", "joint2"], position=[0.1, 0.2])) + world.set_joint_state(ctx, right_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.4])) + updates: list[tuple[str, list[float]]] = [] + shown: list[tuple[str, ...]] = [] + hidden: list[tuple[str, ...]] = [] + sleeps: list[float] = [] + monkeypatch.setattr( + world, + "_set_preview_positions", + lambda _ctx, robot_id, positions: updates.append((robot_id, positions.tolist())), + ) + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: (shown if visible else hidden).append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + monkeypatch.setattr("time.sleep", sleeps.append) + plan = type("Plan", (), {})() + plan.trajectory = _trajectory(["left/joint1", "right/joint2"], [1.0, 2.0], [3.0, 4.0]) + + world.animate_trajectory(plan.trajectory, duration=2.0) + + assert shown == [(left_id,), (right_id,)] + assert hidden == [(left_id,), (right_id,)] + assert updates == [ + (left_id, [1.0, 0.2]), + (right_id, [0.3, 2.0]), + (left_id, [3.0, 0.2]), + (right_id, [0.3, 4.0]), + ] + assert sleeps == [2.0] + + +@requires_drake +def test_drake_animate_trajectory_validates_before_visibility_and_cleans_up( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + shown: list[tuple[str, ...]] = [] + hidden: list[tuple[str, ...]] = [] + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: (shown if visible else hidden).append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + malformed = _trajectory(["unknown/joint1"], [0.0], [1.0]) + with pytest.raises(ValueError, match="unknown robot"): + world.animate_trajectory(malformed) + assert shown == [] + + valid = _trajectory(["arm/joint1"], [0.0], [1.0]) + + def fail_preview_update(*_args: object) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(world, "_set_preview_positions", fail_preview_update) + with pytest.raises(RuntimeError, match="boom"): + world.animate_trajectory(valid) + assert shown == [(robot_id,)] + assert hidden == [(robot_id,)] + + +@requires_drake +def test_drake_cancel_preview_hides_ghosts_before_animation_resumes( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + hidden: list[tuple[str, ...]] = [] + hidden_snapshots_during_sleep: list[list[tuple[str, ...]]] = [] + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: None if visible else hidden.append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + + def cancel_during_sleep(_duration: float) -> None: + world.cancel_preview_animation() + hidden_snapshots_during_sleep.append(list(hidden)) + + monkeypatch.setattr("time.sleep", cancel_during_sleep) + + world.animate_trajectory(_trajectory(["arm/joint1"], [0.0], [1.0])) + + assert hidden_snapshots_during_sleep == [[(robot_id,)]] + assert hidden[0] == (robot_id,) + + +@requires_drake +def test_drake_animate_trajectory_rejects_unknown_robot_before_visibility( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + shown: list[tuple[str, ...]] = [] + with pytest.raises(ValueError, match="unknown robot"): + world.animate_trajectory(_trajectory(["missing/joint1"], [0.0], [1.0])) + + assert shown == [] + + +@requires_drake +def test_drake_animate_trajectory_cancellation_stops_stale_frames_and_hides_preview( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + updates: list[list[float]] = [] + hidden: list[tuple[str, ...]] = [] + monkeypatch.setattr( + world, + "_set_preview_positions", + lambda _ctx, _robot_id, positions: updates.append(positions.tolist()), + ) + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: hidden.append((robot_id,)) if not visible else None, + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + monkeypatch.setattr("time.sleep", lambda _duration: world.cancel_preview_animation()) + + world.animate_trajectory(_trajectory(["arm/joint1"], [1.0], [2.0])) + + assert updates == [[1.0, 0.0]] + assert hidden[0] == (robot_id,) diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py new file mode 100644 index 0000000000..01d52ae89b --- /dev/null +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -0,0 +1,181 @@ +# 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. + +"""Generated-plan materialization tests.""" + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from dimos.manipulation._test_manipulation_helpers import make_module +from dimos.manipulation.manipulation_module import ManipulationState +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import PlanningResult +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +class RecordingGenerator: + calls: list[list[list[float]]] = [] + limits: tuple[list[float], list[float]] | None = None + fail = False + + def __init__( + self, num_joints: int, max_velocity: list[float], max_acceleration: list[float] + ) -> None: + self.num_joints = num_joints + RecordingGenerator.limits = (list(max_velocity), list(max_acceleration)) + + def generate(self, waypoints: list[list[float]]) -> JointTrajectory: + RecordingGenerator.calls.append(waypoints) + if RecordingGenerator.fail: + raise RuntimeError("boom") + return JointTrajectory( + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0] * self.num_joints, + ) + for index, point in enumerate(waypoints) + ] + ) + + +def _robot(name: str, joints: list[str], velocity: float, acceleration: float) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=joints, + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="group", joint_names=tuple(reversed(joints)), base_link="base", tip_link="tip" + ) + ], + max_velocity=velocity, + max_acceleration=acceleration, + ) + + +def _module(monkeypatch: pytest.MonkeyPatch): + RecordingGenerator.calls = [] + RecordingGenerator.limits = None + RecordingGenerator.fail = False + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.JointTrajectoryGenerator", RecordingGenerator + ) + left = _robot("left", ["a", "b"], 1.0, 2.0) + right = _robot("right", ["c"], 3.0, 4.0) + module = make_module() + module._robots = { + "left": ("left_id", left, MagicMock()), + "right": ("right_id", right, MagicMock()), + } + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._planner = MagicMock() + module._state = ManipulationState.PLANNING + module._planning_epoch = 1 + return module + + +def _path(names: list[str], first: list[float], second: list[float]) -> list[JointState]: + return [JointState(name=names, position=first), JointState(name=names, position=second)] + + +def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distinct_path( + monkeypatch, +): + module = _module(monkeypatch) + names = ["left/b", "left/a", "right/c"] + path = _path(names, [0.0, 0.0, 0.0], [0.2, 0.1, 0.3]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + assert module._plan_selected_path(("left/group", "right/group"), path[0], path[-1], 1) + assert RecordingGenerator.calls == [[[0.0, 0.0, 0.0], [0.2, 0.1, 0.3]]] + assert RecordingGenerator.limits == ([1.0, 1.0, 3.0], [2.0, 2.0, 4.0]) + assert module._last_plan is not None + assert module._last_plan.path is not module._last_plan.trajectory.points + assert module._last_plan.trajectory.joint_names == names + assert module._last_plan.trajectory.points[-1].time_from_start == 1.0 + + +@pytest.mark.parametrize( + ("path", "message"), + [ + (_path(["left/a", "left/b"], [0.0, 0.0], [1.0, 1.0]), "joint names"), + (_path(["left/b", "left/a"], [0.0], [1.0]), "dimension"), + (_path(["left/b", "left/a"], [0.0, float("nan")], [1.0, 1.0]), "non-finite"), + ], +) +def test_rejects_malformed_or_nonfinite_waypoints(monkeypatch, path, message): + module = _module(monkeypatch) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) + assert module._last_plan is None + assert message in module._error_message + + +def test_rejects_invalid_limits_and_generator_failure_without_caching(monkeypatch): + module = _module(monkeypatch) + module._robots["left"][1].max_velocity = 0.0 + names = ["left/b", "left/a"] + path = _path(names, [0.0, 0.0], [1.0, 1.0]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) + assert module._last_plan is None + assert RecordingGenerator.calls == [] + + module = _module(monkeypatch) + RecordingGenerator.fail = True + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) + assert module._last_plan is None + assert len(RecordingGenerator.calls) == 1 + + +def test_zero_generation_after_caching_for_status_and_completion(monkeypatch): + module = _module(monkeypatch) + names = ["left/b", "left/a"] + path = _path(names, [0.0, 0.0], [1.0, 1.0]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + assert module._plan_selected_path(("left/group",), path[0], path[-1], 1) + RecordingGenerator.calls = [] + + module._get_coordinator_client = MagicMock(return_value=None) + module._wait_for_trajectory_completion("left", timeout=0.0) + assert RecordingGenerator.calls == [] diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index 069a2adc4c..4a6eb99c63 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -170,13 +170,7 @@ def test_plan_to_joints(self, module, joint_state_zeros): assert module.has_planned_path() is True assert module._last_plan is not None - projected = module._project_plan_to_robot_paths(module._last_plan) - assert projected is not None - path = projected["test_arm"] - assert len(path) > 1 - - traj = module._trajectory_from_robot_path("test_arm", path) - assert traj is not None + traj = module._split_plan_trajectory_by_robot(module._last_plan)["test_arm"] assert len(traj.points) > 1 assert traj.duration > 0 assert module._last_plan.group_ids == ("test_arm/manipulator",) @@ -194,7 +188,7 @@ def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): assert module._last_plan is not None assert module._last_plan.group_ids == ("test_arm/manipulator",) assert module.has_planned_path() is True - assert module._project_plan_to_robot_paths(module._last_plan) is not None + assert module._split_plan_trajectory_by_robot(module._last_plan) is not None def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" @@ -248,11 +242,7 @@ def test_trajectory_name_translation(self, module, joint_state_zeros): assert success is True assert module._last_plan is not None - projected = module._project_plan_to_robot_paths(module._last_plan) - assert projected is not None - path = projected["test_arm"] - traj = module._trajectory_from_robot_path("test_arm", path) - assert traj is not None + traj = module._split_plan_trajectory_by_robot(module._last_plan)["test_arm"] robot_config = module._robots["test_arm"][1] translated = module._translate_trajectory_to_coordinator(traj, robot_config) diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py new file mode 100644 index 0000000000..06cea02e45 --- /dev/null +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -0,0 +1,445 @@ +# 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. + +"""Monitor and preview unit tests for ManipulationModule.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from dimos.manipulation._test_manipulation_helpers import make_module as _make_module +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import GeneratedPlan +from dimos.manipulation.planning.spec.protocols import VisualizationSpec +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +@pytest.fixture +def robot_config_with_mapping() -> RobotModelConfig: + """Create a robot config with joint name mapping.""" + return RobotModelConfig( + name="left_arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["joint1", "joint2", "joint3"], + base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], + joint_name_mapping={ + "left/joint1": "joint1", + "left/joint2": "joint2", + "left/joint3": "joint3", + }, + coordinator_task_name="traj_left", + ) + + +def _one_joint_config(name: str = "arm") -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" + ) + ], + coordinator_task_name=f"traj_{name}", + ) + + +def _install_generated_plan( + module: ManipulationModule, + config: RobotModelConfig, + traj_gen: MagicMock, + *points: list[float], +) -> None: + """Install a generated plan and enough monitor state to derive robot paths.""" + global_joint_names = [f"{config.name}/{joint}" for joint in config.joint_names] + module._robots = {config.name: ("robot_id", config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=config.joint_names, + position=[0.0 for _ in config.joint_names], + ) + module._last_plan = GeneratedPlan( + trajectory=JointTrajectory( + joint_names=global_joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in config.joint_names], + ) + for index, point in enumerate(points) + ], + ), + group_ids=(f"{config.name}/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=global_joint_names, + position=list(point), + ) + for point in points + ], + ) + + +def _make_module_with_monitor(*configs: RobotModelConfig) -> ManipulationModule: + """Create a ManipulationModule with a mocked world monitor and robots configured.""" + module = _make_module() + module._world_monitor = MagicMock() + module._init_joints = {} + for config in configs: + robot_id = f"robot_{config.name}" + module._robots[config.name] = (robot_id, config, MagicMock()) + return module + + +def _make_joint_state(positions: list[float], name: list[str] | None = None) -> JointState: + return JointState(name=name or [f"j{i}" for i in range(len(positions))], position=positions) + + +def _make_path(*points: list[float]) -> list[JointState]: + return [_make_joint_state(list(point)) for point in points] + + +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( + joint_names=joint_names, + points=[ + TrajectoryPoint(time_from_start=time_from_start, positions=positions) + for time_from_start, positions in points + ], + ) + + +def _make_world_monitor_with_viz(viz: VisualizationSpec | None) -> WorldMonitor: + world = MagicMock() + return WorldMonitor( + world=world, + visualization=viz, + ) + + +class FakeVisualization: + def __init__(self) -> None: + self.close_count = 0 + self.published = False + self.preview_shown: list[str] = [] + self.preview_hidden: list[str] = [] + self.animations: list[tuple[str, list[JointState], float]] = [] + self.preview_animation_cancellations = 0 + + def initialize(self, session) -> None: + pass + + def get_visualization_url(self) -> str | None: + return "123" + + def update_state(self, frame) -> None: + self.published = True + + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: + self.animations.append( + ( + tuple(trajectory.joint_names), + list(trajectory.points), + duration if duration is not None else 0.0, + ) + ) + + def cancel_preview_animation(self) -> None: + self.preview_animation_cancellations += 1 + + def close(self) -> None: + self.close_count += 1 + + +class TestOnJointState: + """Test _on_joint_state routing, splitting, and init capture.""" + + def test_routes_positions_to_monitor(self, robot_config_with_mapping): + """Joint positions from aggregated message are routed to the correct monitor.""" + module = _make_module_with_monitor(robot_config_with_mapping) + + msg = JointState( + name=["left/joint1", "left/joint2", "left/joint3"], + position=[0.1, 0.2, 0.3], + velocity=[1.0, 2.0, 3.0], + ) + module._on_joint_state(msg) + + # Verify world_monitor received the sub-message + module._world_monitor.on_joint_state.assert_called_once() + call_args = module._world_monitor.on_joint_state.call_args + sub_msg = call_args[0][0] + assert sub_msg.position == [0.1, 0.2, 0.3] + assert sub_msg.velocity == [1.0, 2.0, 3.0] + assert call_args[1]["robot_id"] == "robot_left_arm" + + def test_skips_robot_with_missing_joints(self, robot_config_with_mapping): + """Robots whose joints are absent from the message are skipped.""" + module = _make_module_with_monitor(robot_config_with_mapping) + + # Message has none of left_arm's joints + msg = JointState( + name=["right/joint1", "right/joint2"], + position=[0.5, 0.6], + ) + module._on_joint_state(msg) + + module._world_monitor.on_joint_state.assert_not_called() + + def test_captures_init_joints_on_first_call(self, robot_config_with_mapping): + """First joint state is stored as init joints; subsequent calls don't overwrite.""" + module = _make_module_with_monitor(robot_config_with_mapping) + + first_msg = JointState( + name=["left/joint1", "left/joint2", "left/joint3"], + position=[0.1, 0.2, 0.3], + ) + module._on_joint_state(first_msg) + assert "left_arm" in module._init_joints + assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] + + # Second call should NOT overwrite + second_msg = JointState( + name=["left/joint1", "left/joint2", "left/joint3"], + position=[0.9, 0.8, 0.7], + ) + module._on_joint_state(second_msg) + assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] + + def test_multi_robot_splits_correctly(self): + """With two robots, each gets only its own joints from the aggregated message.""" + left_config = RobotModelConfig( + name="left", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j1", "j2"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], + joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, + coordinator_task_name="traj_left", + ) + right_config = RobotModelConfig( + name="right", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j1", "j2"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], + joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, + coordinator_task_name="traj_right", + ) + module = _make_module_with_monitor(left_config, right_config) + + msg = JointState( + name=["left/j1", "left/j2", "right/j1", "right/j2"], + position=[1.0, 2.0, 3.0, 4.0], + velocity=[0.1, 0.2, 0.3, 0.4], + ) + module._on_joint_state(msg) + + assert module._world_monitor.on_joint_state.call_count == 2 + + # Collect calls by robot_id + calls = { + call[1]["robot_id"]: call[0][0] + for call in module._world_monitor.on_joint_state.call_args_list + } + assert calls["robot_left"].position == [1.0, 2.0] + assert calls["robot_right"].position == [3.0, 4.0] + assert calls["robot_left"].velocity == [0.1, 0.2] + assert calls["robot_right"].velocity == [0.3, 0.4] + + def test_no_monitor_returns_early(self, robot_config_with_mapping): + """When world_monitor is None, _on_joint_state returns without error.""" + module = _make_module() + module._robots = {"left_arm": ("id", robot_config_with_mapping, MagicMock())} + module._world_monitor = None + + # Should not raise + msg = JointState( + name=["left/joint1", "left/joint2", "left/joint3"], + position=[0.1, 0.2, 0.3], + ) + module._on_joint_state(msg) + + +class TestWorldMonitorVisualization: + def test_visualization_routing_and_stop_all_monitors(self): + viz = FakeVisualization() + monitor = _make_world_monitor_with_viz(viz) + state_monitor = MagicMock() + obstacle_monitor = MagicMock() + monitor._state_monitors = {"robot": state_monitor} + monitor._obstacle_monitor = obstacle_monitor + monitor._viz_thread = MagicMock() + monitor._viz_thread.is_alive.return_value = False + + assert monitor.get_visualization_url() == "123" + monitor.update_visualization_state() + monitor.cancel_preview_animation() + path = _make_path([1.0], [2.0], [3.0]) + plan = GeneratedPlan( + trajectory=JointTrajectory(), + group_ids=("robot/group",), + status=PlanningStatus.SUCCESS, + path=path, + ) + monitor.animate_trajectory(plan.trajectory, 4.5) + assert monitor.visualization is viz + assert viz.published is True + assert viz.preview_animation_cancellations == 2 + assert viz.animations == [(tuple(), [], 4.5)] + + monitor.stop_all_monitors() + + assert viz.close_count == 1 + state_monitor.stop.assert_called_once() + obstacle_monitor.stop.assert_called_once() + + def test_visualization_none_is_noop(self): + monitor = _make_world_monitor_with_viz(None) + + assert monitor.get_visualization_url() is None + monitor.update_visualization_state() + monitor.cancel_preview_animation() + monitor.animate_trajectory(JointTrajectory(), 1.0) + monitor.start_visualization_thread() + assert monitor._viz_thread is None + + +class TestManipulationPreview: + def test_clear_planned_path_invalidates_before_dismissing_preview(self): + module = _make_module() + plan = GeneratedPlan(trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[]) + module._last_plan = plan + module._world_monitor = MagicMock() + plan_during_dismissal: list[GeneratedPlan | None] = [] + module._world_monitor.cancel_preview_animation.side_effect = ( + lambda: plan_during_dismissal.append(module._last_plan) + ) + + assert module.clear_planned_path() is True + + assert plan_during_dismissal == [None] + module._world_monitor.cancel_preview_animation.assert_called_once_with() + assert module._last_plan is None + + def test_clear_planned_path_clears_without_a_world_monitor(self): + module = _make_module() + module._last_plan = GeneratedPlan( + trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] + ) + + assert module.clear_planned_path() is True + assert module._last_plan is None + + def test_dismiss_preview_noop_without_monitor(self): + module = _make_module() + + module._dismiss_preview(["arm/manipulator"]) + + def test_dismiss_preview_routes_to_monitor(self): + module = _make_module() + module._world_monitor = MagicMock() + + module._dismiss_preview(["arm/manipulator"]) + + module._world_monitor.cancel_preview_animation.assert_called_once_with() + + def test_preview_routes_one_complete_plan_with_default_duration(self): + module = _make_module() + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0], [2.0]) + + assert module.preview_plan() is True + + module._world_monitor.animate_trajectory.assert_called_once_with( + module._last_plan.trajectory, None + ) + + def test_preview_robot_name_validates_affectedness_without_trimming(self): + module = _make_module() + left = _one_joint_config("left") + right = _one_joint_config("right") + traj_gen = MagicMock() + module._robots = { + "left": ("left_id", left, traj_gen), + "right": ("right_id", right, traj_gen), + } + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._last_plan = GeneratedPlan( + trajectory=JointTrajectory(), + group_ids=("left/manipulator", "right/manipulator"), + status=PlanningStatus.SUCCESS, + path=[ + JointState(name=["left/j0", "right/j0"], position=[0.0, 0.0]), + JointState(name=["left/j0", "right/j0"], position=[1.0, 1.0]), + ], + ) + + assert module.preview_plan(duration=2.5, robot_name="left") is True + + module._world_monitor.animate_trajectory.assert_called_once_with( + module._last_plan.trajectory, 2.5 + ) + + def test_preview_rejects_unaffected_compatibility_robot(self): + module = _make_module() + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) + + assert module.preview_plan(robot_name="other") is False + module._world_monitor.animate_trajectory.assert_not_called() diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 3e27bfef7a..4f37c28fa0 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -17,12 +17,12 @@ from __future__ import annotations from pathlib import Path -import threading from unittest.mock import MagicMock import pytest from pytest_mock import MockerFixture +from dimos.manipulation._test_manipulation_helpers import make_module as _make_module from dimos.manipulation.manipulation_module import ( ManipulationModule, ManipulationModuleConfig, @@ -38,9 +38,7 @@ GeneratedPlan, IKResult, PlanningResult, - PlanningSceneInfo, ) -from dimos.manipulation.planning.spec.protocols import VisualizationSpec from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -115,26 +113,6 @@ def simple_trajectory(): ) -class _ManipulationModuleHarness(ManipulationModule): - def __init__(self) -> None: - self._state = ManipulationState.IDLE - self._lock = threading.Lock() - self._error_message = "" - self._planning_epoch = 0 - self._robots = {} - self._last_plan = None - self._world_monitor = None - self._planner = None - self._kinematics = None - self._coordinator_client = None - self.config = MagicMock(planning_timeout=10.0) - - -def _make_module() -> ManipulationModule: - """Create a lightweight ManipulationModule harness for behavior tests.""" - return _ManipulationModuleHarness() - - def _one_joint_config(name: str = "arm") -> RobotModelConfig: return RobotModelConfig( name=name, @@ -158,6 +136,7 @@ def _install_generated_plan( *points: list[float], ) -> None: """Install a generated plan and enough monitor state to derive robot paths.""" + global_joint_names = [f"{config.name}/{joint}" for joint in config.joint_names] module._robots = {config.name: ("robot_id", config, traj_gen)} module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([config]) @@ -165,12 +144,27 @@ def _install_generated_plan( name=config.joint_names, position=[0.0 for _ in config.joint_names], ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=global_joint_names, + position=[0.0 for _ in config.joint_names], + ) module._last_plan = GeneratedPlan( + trajectory=JointTrajectory( + joint_names=global_joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in config.joint_names], + ) + for index, point in enumerate(points) + ], + ), group_ids=(f"{config.name}/manipulator",), status=PlanningStatus.SUCCESS, path=[ JointState( - name=[f"{config.name}/{joint}" for joint in config.joint_names], + name=global_joint_names, position=list(point), ) for point in points @@ -178,6 +172,31 @@ def _install_generated_plan( ) +def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> JointTrajectory: + return JointTrajectory( + joint_names=joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in joint_names], + ) + for index, point in enumerate(points) + ], + ) + + +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( + joint_names=joint_names, + points=[ + TrajectoryPoint(time_from_start=time_from_start, positions=positions) + for time_from_start, positions in points + ], + ) + + class TestStateMachine: """Test state transitions.""" @@ -197,6 +216,18 @@ def test_cancel_interrupts_active_work(self): assert module.cancel() is True assert module._state == ManipulationState.IDLE + def test_cancel_hides_active_plan_preview(self): + module = _make_module() + module._state = ManipulationState.EXECUTING + module._last_plan = GeneratedPlan( + trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] + ) + module._world_monitor = MagicMock() + + assert module.cancel() is True + + module._world_monitor.cancel_preview_animation.assert_called_once_with() + def test_reset_not_during_execution(self): """Reset works in any state except EXECUTING.""" module = _make_module() @@ -485,6 +516,14 @@ def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches(self, rob name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) result_path = [ JointState( name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], @@ -518,10 +557,9 @@ def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches(self, rob assert module._last_plan is not None assert module._last_plan.group_ids == ("test_arm/manipulator",) assert module._last_plan.path == result_path - projected = module._project_plan_to_robot_paths(module._last_plan) - assert projected is not None - assert projected["test_arm"][1].position == [0.1, 0.2, 0.3] - assert module._trajectory_from_robot_path("test_arm", projected["test_arm"]) is not None + trajectories = module._split_plan_trajectory_by_robot(module._last_plan) + assert trajectories is not None + assert trajectories["test_arm"].points[-1].positions == [0.1, 0.2, 0.3] module._planner.plan_selected_joint_path.assert_called_once() _, kwargs = module._planner.plan_selected_joint_path.call_args assert kwargs["selection"].group_ids == ("test_arm/manipulator",) @@ -531,6 +569,18 @@ def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches(self, rob "test_arm/joint3", ] + success = module.plan_to_joint_targets( + { + "test_arm/manipulator": JointState( + name=robot_config.joint_names, + position=[0.1, 0.2, 0.3], + ) + } + ) + + assert success is True + assert module._planner.plan_selected_joint_path.call_count == 2 + def test_plan_to_pose_targets_uses_group_ik_and_selected_path(self, robot_config): module = _make_module() registry = PlanningGroupRegistry([robot_config]) @@ -550,6 +600,10 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path(self, robot_config name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) ik_goal = JointState( name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], position=[0.1, 0.2, 0.3], @@ -589,7 +643,9 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path(self, robot_config _, planner_kwargs = module._planner.plan_selected_joint_path.call_args assert planner_kwargs["goal"] is ik_goal - def test_failed_plan_projection_clears_generated_plan(self, robot_config, simple_trajectory): + def test_failed_plan_materialization_clears_generated_plan( + self, robot_config, simple_trajectory + ): module = _make_module() registry = PlanningGroupRegistry([robot_config]) module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} @@ -602,6 +658,11 @@ def test_failed_plan_projection_clears_generated_plan(self, robot_config, simple ) module._world_monitor.get_current_joint_state.return_value = None module._last_plan = GeneratedPlan( + trajectory=_generated_plan_trajectory( + ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + [0.0, 0.0, 0.0], + [0.1, 0.2, 0.3], + ), group_ids=("test_arm/manipulator",), status=PlanningStatus.SUCCESS, path=[ @@ -647,7 +708,16 @@ def test_execute_plan_dispatches_each_affected_robot(self, robot_config, simple_ name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) module._last_plan = GeneratedPlan( + trajectory=_generated_plan_trajectory( + ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + [0.0, 0.0, 0.0], + [0.1, 0.2, 0.3], + ), group_ids=("test_arm/manipulator",), status=PlanningStatus.SUCCESS, path=[ @@ -673,8 +743,165 @@ def test_execute_plan_dispatches_each_affected_robot(self, robot_config, simple_ assert method_name == "execute" trajectory = payload["trajectory"] assert trajectory.joint_names == simple_trajectory.joint_names - assert trajectory.points == simple_trajectory.points + assert [point.positions for point in trajectory.points] == [ + [0.0, 0.0, 0.0], + [0.1, 0.2, 0.3], + ] + traj_gen.generate.assert_not_called() + assert module._state == ManipulationState.COMPLETED + + def test_execute_plan_rejects_stale_planned_robot_without_consuming_plan(self): + module = _make_module() + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) + module._state = ManipulationState.COMPLETED + module._world_monitor.current_global_joint_state.return_value = JointState( + name=[], position=[] + ) + module._world_monitor.is_state_stale.return_value = True + mock_client = MagicMock() + module._coordinator_client = mock_client + + assert module.execute_plan() is False + + assert module._last_plan is not None assert module._state == ManipulationState.COMPLETED + mock_client.task_invoke.assert_not_called() + + def test_execute_plan_rejects_duplicate_current_global_joints_without_consuming_plan(self): + module = _make_module() + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) + module._state = ManipulationState.COMPLETED + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["arm/j0", "arm/j0"], position=[0.0, 0.0] + ) + mock_client = MagicMock() + module._coordinator_client = mock_client + + assert module.execute_plan() is False + + assert module._last_plan is not None + assert module._state == ManipulationState.COMPLETED + mock_client.task_invoke.assert_not_called() + + def test_execute_plan_rejects_reordered_current_global_joints_without_consuming_plan(self): + module = _make_module() + config = RobotModelConfig( + name="arm", + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0", "j1"], + base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j0", "j1"), + base_link="base_link", + tip_link="ee", + ) + ], + coordinator_task_name="traj_arm", + ) + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0, 1.0], [0.5, 1.5]) + module._state = ManipulationState.COMPLETED + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["arm/j1", "arm/j0"], position=[1.0, 0.0] + ) + mock_client = MagicMock() + module._coordinator_client = mock_client + + assert module.execute_plan() is False + + assert module._last_plan is not None + assert module._state == ManipulationState.COMPLETED + mock_client.task_invoke.assert_not_called() + + def test_execute_plan_dispatches_selected_subsets_with_shared_clock_and_mapping(self): + left = RobotModelConfig( + name="left", + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0", "j1"], + planning_groups=[ + PlanningGroupDefinition( + name="wrist", joint_names=("j1",), base_link="base", tip_link="ee" + ) + ], + joint_name_mapping={"left_coord_j1": "j1"}, + coordinator_task_name="traj_left", + ) + right = RobotModelConfig( + name="right", + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["k0", "k1"], + planning_groups=[ + PlanningGroupDefinition( + name="elbow", joint_names=("k0",), base_link="base", tip_link="ee" + ) + ], + coordinator_task_name="traj_right", + ) + module = _make_module() + module._robots = { + "left": ("left_id", left, MagicMock()), + "right": ("right_id", right, MagicMock()), + } + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._last_plan = GeneratedPlan( + group_ids=("left/wrist", "right/elbow"), + status=PlanningStatus.SUCCESS, + path=[ + JointState(name=["left/j1", "right/k0"], position=[0.0, 1.0]), + JointState(name=["left/j1", "right/k0"], position=[0.5, 1.5]), + ], + trajectory=JointTrajectory( + joint_names=["left/j1", "right/k0"], + points=[ + TrajectoryPoint( + time_from_start=0.0, + positions=[0.0, 1.0], + velocities=[0.0, 0.0], + ), + TrajectoryPoint( + time_from_start=2.5, + positions=[0.5, 1.5], + velocities=[0.2, 0.4], + ), + ], + ), + ) + mock_client = MagicMock() + mock_client.task_invoke.return_value = True + module._coordinator_client = mock_client + module._world_monitor.get_current_joint_state.side_effect = lambda robot_id: { + "left_id": JointState(name=["j0", "j1"], position=[9.0, 0.0]), + "right_id": JointState(name=["k0", "k1"], position=[1.0, 9.0]), + }[robot_id] + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["left/j1", "right/k0"], position=[0.0, 1.0] + ) + + assert module.execute_plan() is True + + calls = mock_client.task_invoke.call_args_list + left_payload = calls[0].args[2]["trajectory"] + right_payload = calls[1].args[2]["trajectory"] + assert calls[0].args[0] == "traj_left" + assert left_payload.joint_names == ["left_coord_j1"] + assert [point.time_from_start for point in left_payload.points] == [0.0, 2.5] + assert [point.positions for point in left_payload.points] == [[0.0], [0.5]] + assert [point.velocities for point in left_payload.points] == [[0.0], [0.2]] + assert calls[1].args[0] == "traj_right" + assert right_payload.joint_names == ["k0"] + assert [point.time_from_start for point in right_payload.points] == [0.0, 2.5] + assert [point.positions for point in right_payload.points] == [[1.0], [1.5]] + assert [point.velocities for point in right_payload.points] == [[0.0], [0.4]] def test_pose_wrappers_fail_safely_without_unique_pose_group(self, robot_config): no_pose_config = RobotModelConfig( @@ -819,8 +1046,17 @@ def test_execute_plan_requires_task_name_without_sticking_executing(self, robot_ name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) module._coordinator_client = MagicMock() module._last_plan = GeneratedPlan( + trajectory=_generated_plan_trajectory( + ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + [0.0, 0.0, 0.0], + [0.1, 0.2, 0.3], + ), group_ids=("test_arm/manipulator",), status=PlanningStatus.SUCCESS, path=[ @@ -836,7 +1072,7 @@ def test_execute_plan_requires_task_name_without_sticking_executing(self, robot_ ) assert module.execute_plan() is False - assert module._state == ManipulationState.FAULT + assert module._state == ManipulationState.IDLE def test_execute_success(self, robot_config, simple_trajectory): """Successful execute calls coordinator via task_invoke.""" @@ -857,7 +1093,11 @@ def test_execute_success(self, robot_config, simple_trajectory): assert method_name == "execute" trajectory = payload["trajectory"] assert trajectory.joint_names == simple_trajectory.joint_names - assert trajectory.points == simple_trajectory.points + assert [point.positions for point in trajectory.points] == [ + [0.0, 0.0, 0.0], + [0.1, 0.2, 0.3], + ] + traj_gen.generate.assert_not_called() def test_execute_rejected(self, robot_config, simple_trajectory): """Rejected execution sets FAULT state.""" @@ -888,321 +1128,3 @@ def test_bidirectional_mapping(self, robot_config_with_mapping): # URDF -> Coordinator assert config.get_coordinator_joint_name("joint1") == "left/joint1" assert config.get_coordinator_joint_name("unknown") == "unknown" - - -def _make_module_with_monitor(*configs: RobotModelConfig) -> ManipulationModule: - """Create a ManipulationModule with a mocked world monitor and robots configured.""" - module = _make_module() - module._world_monitor = MagicMock() - module._init_joints = {} - for config in configs: - robot_id = f"robot_{config.name}" - module._robots[config.name] = (robot_id, config, MagicMock()) - return module - - -def _make_joint_state(positions: list[float], name: list[str] | None = None) -> JointState: - return JointState(name=name or [f"j{i}" for i in range(len(positions))], position=positions) - - -def _make_path(*points: list[float]) -> list[JointState]: - return [_make_joint_state(list(point)) for point in points] - - -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( - joint_names=joint_names, - points=[ - TrajectoryPoint(time_from_start=time_from_start, positions=positions) - for time_from_start, positions in points - ], - ) - - -def _make_world_monitor_with_viz(viz: VisualizationSpec | None) -> WorldMonitor: - world = MagicMock() - return WorldMonitor( - world=world, - visualization=viz, - ) - - -class FakeVisualization: - def __init__(self) -> None: - self.close_count = 0 - self.published = False - self.preview_shown: list[str] = [] - self.preview_hidden: list[str] = [] - self.animations: list[tuple[str, list[JointState], float]] = [] - - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - pass - - def get_visualization_url(self) -> str | None: - return "123" - - def publish_visualization(self, ctx: object | None = None) -> None: - self.published = True - - def show_preview(self, robot_id: str) -> None: - self.preview_shown.append(robot_id) - - def hide_preview(self, robot_id: str) -> None: - self.preview_hidden.append(robot_id) - - def animate_path(self, robot_id: str, path: list[JointState], duration: float = 3.0) -> None: - self.animations.append((robot_id, path, duration)) - - def close(self) -> None: - self.close_count += 1 - - -class TestOnJointState: - """Test _on_joint_state routing, splitting, and init capture.""" - - def test_routes_positions_to_monitor(self, robot_config_with_mapping): - """Joint positions from aggregated message are routed to the correct monitor.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - velocity=[1.0, 2.0, 3.0], - ) - module._on_joint_state(msg) - - # Verify world_monitor received the sub-message - module._world_monitor.on_joint_state.assert_called_once() - call_args = module._world_monitor.on_joint_state.call_args - sub_msg = call_args[0][0] - assert sub_msg.position == [0.1, 0.2, 0.3] - assert sub_msg.velocity == [1.0, 2.0, 3.0] - assert call_args[1]["robot_id"] == "robot_left_arm" - - def test_skips_robot_with_missing_joints(self, robot_config_with_mapping): - """Robots whose joints are absent from the message are skipped.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - # Message has none of left_arm's joints - msg = JointState( - name=["right/joint1", "right/joint2"], - position=[0.5, 0.6], - ) - module._on_joint_state(msg) - - module._world_monitor.on_joint_state.assert_not_called() - - def test_captures_init_joints_on_first_call(self, robot_config_with_mapping): - """First joint state is stored as init joints; subsequent calls don't overwrite.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - first_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(first_msg) - assert "left_arm" in module._init_joints - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - # Second call should NOT overwrite - second_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.9, 0.8, 0.7], - ) - module._on_joint_state(second_msg) - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - def test_multi_robot_splits_correctly(self): - """With two robots, each gets only its own joints from the aggregated message.""" - left_config = RobotModelConfig( - name="left", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, - coordinator_task_name="traj_left", - ) - right_config = RobotModelConfig( - name="right", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, - coordinator_task_name="traj_right", - ) - module = _make_module_with_monitor(left_config, right_config) - - msg = JointState( - name=["left/j1", "left/j2", "right/j1", "right/j2"], - position=[1.0, 2.0, 3.0, 4.0], - velocity=[0.1, 0.2, 0.3, 0.4], - ) - module._on_joint_state(msg) - - assert module._world_monitor.on_joint_state.call_count == 2 - - # Collect calls by robot_id - calls = { - call[1]["robot_id"]: call[0][0] - for call in module._world_monitor.on_joint_state.call_args_list - } - assert calls["robot_left"].position == [1.0, 2.0] - assert calls["robot_right"].position == [3.0, 4.0] - assert calls["robot_left"].velocity == [0.1, 0.2] - assert calls["robot_right"].velocity == [0.3, 0.4] - - def test_no_monitor_returns_early(self, robot_config_with_mapping): - """When world_monitor is None, _on_joint_state returns without error.""" - module = _make_module() - module._robots = {"left_arm": ("id", robot_config_with_mapping, MagicMock())} - module._world_monitor = None - - # Should not raise - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(msg) - - -class TestWorldMonitorVisualization: - def test_visualization_routing_and_stop_all_monitors(self): - viz = FakeVisualization() - monitor = _make_world_monitor_with_viz(viz) - state_monitor = MagicMock() - obstacle_monitor = MagicMock() - monitor._state_monitors = {"robot": state_monitor} - monitor._obstacle_monitor = obstacle_monitor - monitor._viz_thread = MagicMock() - monitor._viz_thread.is_alive.return_value = False - - assert monitor.get_visualization_url() == "123" - monitor.publish_visualization() - monitor.show_preview("robot") - monitor.hide_preview("robot") - path = _make_path([1.0], [2.0], [3.0]) - monitor.animate_path("robot", path, 4.5) - assert monitor.visualization is viz - assert viz.published is True - assert viz.preview_shown == ["robot"] - assert viz.preview_hidden == ["robot"] - assert viz.animations == [("robot", path, 4.5)] - - monitor.stop_all_monitors() - - assert viz.close_count == 1 - state_monitor.stop.assert_called_once() - obstacle_monitor.stop.assert_called_once() - - def test_visualization_none_is_noop(self): - monitor = _make_world_monitor_with_viz(None) - - assert monitor.get_visualization_url() is None - monitor.publish_visualization() - monitor.show_preview("robot") - monitor.hide_preview("robot") - monitor.animate_path("robot", [1], 1.0) - monitor.start_visualization_thread() - assert monitor._viz_thread is None - - -class TestManipulationPreview: - def test_dismiss_preview_noop_without_monitor(self): - module = _make_module() - - module._dismiss_preview("robot_id") - - def test_dismiss_preview_routes_to_monitor(self): - module = _make_module() - module._world_monitor = MagicMock() - - module._dismiss_preview("robot_id") - - module._world_monitor.hide_preview.assert_called_once_with("robot_id") - module._world_monitor.publish_visualization.assert_called_once_with() - - def test_preview_path_uses_trajectory_duration_and_interpolates(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (2.0, [2.0])) - _install_generated_plan(module, config, traj_gen, [0.0], [2.0]) - - assert module.preview_path(robot_name="arm", target_fps=2.0) is True - - module._world_monitor.animate_path.assert_called_once() - robot_id, preview_path, duration = module._world_monitor.animate_path.call_args.args - assert robot_id == "robot_id" - assert duration == 2.0 - assert [state.position for state in preview_path] == [[0.0], [0.5], [1.0], [1.5], [2.0]] - - def test_preview_path_explicit_duration_overrides_and_fps_densifies(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (9.0, [9.0])) - _install_generated_plan(module, config, traj_gen, [0.0], [9.0]) - - assert module.preview_path(duration=1.5, robot_name="arm", target_fps=2.0) is True - - module._world_monitor.animate_path.assert_called_once() - robot_id, preview_path, duration = module._world_monitor.animate_path.call_args.args - assert robot_id == "robot_id" - assert duration == 1.5 - assert [state.position for state in preview_path] == [[0.0], [3.0], [6.0], [9.0]] - - def test_preview_path_returns_false_when_trajectory_cannot_be_derived(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - _install_generated_plan(module, config, traj_gen, [0.0]) - - assert module.preview_path(robot_name="arm", target_fps=10.0) is False - module._world_monitor.animate_path.assert_not_called() - - def test_preview_path_skips_interpolation_for_nonpositive_fps_or_duration(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (2.0, [1.0])) - _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) - - assert module.preview_path(robot_name="arm", target_fps=0.0) is True - assert module.preview_path(duration=0.0, robot_name="arm", target_fps=20.0) is True - - assert [ - state.position for state in module._world_monitor.animate_path.call_args_list[0].args[1] - ] == [[0.0], [1.0]] - assert [ - state.position for state in module._world_monitor.animate_path.call_args_list[1].args[1] - ] == [[0.0], [1.0]] - - def test_preview_path_returns_false_for_missing_inputs(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (1.0, [1.0])) - _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) - module._world_monitor = None - - assert module.preview_path(robot_name="arm") is False - - module._world_monitor = MagicMock() - module._robots = {} - assert module.preview_path(robot_name="arm") is False - - _install_generated_plan(module, config, traj_gen) - assert module.preview_path(robot_name="arm") is False diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py new file mode 100644 index 0000000000..74eee2cd1a --- /dev/null +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -0,0 +1,271 @@ +# 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. + +"""Tests for planned-path execution reservation and result propagation.""" + +from collections.abc import Callable +from pathlib import Path +import threading +import time +from unittest.mock import MagicMock + +from dimos.manipulation._test_manipulation_helpers import make_module +from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationState +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus +from dimos.manipulation.planning.spec.models import GeneratedPlan, IKResult, PlanningResult +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +def _plan(group_id: str = "arm/manipulator") -> GeneratedPlan: + return GeneratedPlan( + group_ids=(group_id,), + path=[ + JointState(name=["arm/j0"], position=[0.0]), + JointState(name=["arm/j0"], position=[1.0]), + ], + trajectory=JointTrajectory( + joint_names=["arm/j0"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), + TrajectoryPoint(time_from_start=1.0, positions=[1.0], velocities=[0.0]), + ], + ), + ) + + +def _module_with_current(current: JointState) -> ManipulationModule: + module = make_module() + robot_config = RobotModelConfig( + name="arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j0",), + base_link="base", + tip_link="tool", + ) + ], + max_velocity=1.0, + max_acceleration=1.0, + coordinator_task_name="traj_arm", + ) + module._robots = {"arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) + module._world_monitor.current_global_joint_state.return_value = current + module._world_monitor.is_state_stale.return_value = False + module._coordinator_client = MagicMock() + module._coordinator_client.task_invoke.return_value = True + return module + + +def _install_current_global_state(module: ManipulationModule, current: JointState) -> None: + robot_config = RobotModelConfig( + name="arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base", tip_link="tool" + ) + ], + coordinator_task_name="traj_arm", + ) + module._robots = {"arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) + module._world_monitor.current_global_joint_state.return_value = current + module._coordinator_client = MagicMock() + module._coordinator_client.task_invoke.return_value = True + + +def test_plan_to_pose_targets_propagates_failed_plan_as_false(): + module = make_module() + robot_config = RobotModelConfig( + name="test_arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["joint1", "joint2", "joint3"], + base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], + max_velocity=1.0, + max_acceleration=2.0, + coordinator_task_name="traj_arm", + ) + registry = PlanningGroupRegistry([robot_config]) + module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = registry + module._world_monitor.current_global_joint_state.return_value = JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.0, 0.0, 0.0], + ) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=robot_config.joint_names, + position=[0.0, 0.0, 0.0], + ) + module._kinematics = MagicMock() + module._kinematics.solve_pose_targets.return_value = IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + position=[0.1, 0.2, 0.3], + ), + ) + module._planner = MagicMock() + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.NO_SOLUTION + ) + pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) + + assert module.plan_to_pose_targets({"test_arm/manipulator": pose}) is False + assert module._state == ManipulationState.FAULT + + +def test_execute_plan_selects_stored_plan_under_gate_and_reuses_it(): + module = make_module() + plan = _plan() + module._last_plan = plan + module._state = ManipulationState.COMPLETED + _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) + + assert module.execute_plan() is True + assert module.execute_plan() is True + assert module._coordinator_client.task_invoke.call_count == 2 + + +def test_execute_and_execute_plan_race_without_consuming_latest_plan(): + module = make_module() + plan = _plan() + module._last_plan = plan + module._state = ManipulationState.COMPLETED + _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) + callers_ready = threading.Barrier(3) + dispatch_started = threading.Event() + allow_dispatch_return = threading.Event() + results: list[bool] = [] + + def dispatch(*_args, **_kwargs): + dispatch_started.set() + assert allow_dispatch_return.wait(timeout=1.0) + return True + + def call_execute(execute: Callable[[], bool]) -> None: + callers_ready.wait(timeout=1.0) + results.append(execute()) + + module._coordinator_client.task_invoke.side_effect = dispatch + legacy = threading.Thread(target=call_execute, args=(module.execute,)) + current = threading.Thread(target=call_execute, args=(module.execute_plan,)) + legacy.start() + current.start() + callers_ready.wait(timeout=1.0) + assert dispatch_started.wait(timeout=1.0) + time.sleep(0.05) + + allow_dispatch_return.set() + legacy.join(timeout=1.0) + current.join(timeout=1.0) + + assert not legacy.is_alive() + assert not current.is_alive() + assert sorted(results) == [False, True] + assert module._coordinator_client.task_invoke.call_count == 1 + + +def test_execute_plan_accepts_a_direct_plan_without_reserving_stored_plan(): + module = make_module() + direct_plan = _plan() + module._last_plan = _plan("other/manipulator") + _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) + + assert module.execute_plan(plan=direct_plan) is True + module._coordinator_client.task_invoke.assert_called_once() + + +def test_execute_plan_uses_current_global_state_without_consuming_cached_plan(): + module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) + cached_plan = _plan("cached/manipulator") + module._last_plan = cached_plan + assert module.execute_plan(plan=_plan()) is True + module._coordinator_client.task_invoke.assert_called_once() + assert module._last_plan is cached_plan + + +def test_execute_plan_rejects_missing_direct_plan_start_before_dispatch_without_reserving_cached_plan(): + module = _module_with_current(JointState(name=[], position=[])) + cached_plan = _plan("cached/manipulator") + module._last_plan = cached_plan + + assert module.execute_plan(plan=_plan()) is False + module._coordinator_client.task_invoke.assert_not_called() + assert module._last_plan is cached_plan + + +def test_execute_plan_rejects_malformed_direct_plan_before_dispatch_without_reserving_cached_plan(): + module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) + cached_plan = _plan("cached/manipulator") + direct_plan = _plan() + direct_plan.trajectory.points[0].positions = [] + module._last_plan = cached_plan + + assert module.execute_plan(plan=direct_plan) is False + module._coordinator_client.task_invoke.assert_not_called() + assert module._last_plan is cached_plan + + +def test_execute_plan_pre_dispatch_exception_restores_previous_state_without_faulting(): + module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) + module._last_plan = _plan() + module._state = ManipulationState.COMPLETED + module._prepare_execution = MagicMock(side_effect=RuntimeError("split exploded")) + + assert module.execute_plan() is False + + assert module._state == ManipulationState.COMPLETED + assert module._error_message == "Failed to prepare execution: split exploded" + module._coordinator_client.task_invoke.assert_not_called() + + +def test_execute_plan_dispatch_exception_faults_without_sticking_executing(): + module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) + module._last_plan = _plan() + module._coordinator_client.task_invoke.side_effect = RuntimeError("rpc exploded") + + assert module.execute_plan() is False + + assert module._state == ManipulationState.FAULT + assert module._error_message == "Failed to dispatch trajectory: rpc exploded" diff --git a/dimos/manipulation/visualization/factory.py b/dimos/manipulation/visualization/factory.py index 3e4e5c25f0..cbb2fb397d 100644 --- a/dimos/manipulation/visualization/factory.py +++ b/dimos/manipulation/visualization/factory.py @@ -53,10 +53,7 @@ def create_manipulation_visualization( ViserManipulationVisualizer, ) - return ViserManipulationVisualizer( - world_monitor=world_monitor, - manipulation_module=manipulation_module, - config=config, - ) + _ = world_monitor, manipulation_module + return ViserManipulationVisualizer(config=config) raise AssertionError(f"Unhandled manipulation visualization config: {config!r}") diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py new file mode 100644 index 0000000000..21af8dfd1b --- /dev/null +++ b/dimos/manipulation/visualization/operator.py @@ -0,0 +1,343 @@ +# 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. + +"""UI-neutral facade for interactive manipulation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +import math +from typing import TYPE_CHECKING, cast + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID, RobotName +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + +if TYPE_CHECKING: + from dimos.manipulation.manipulation_module import ManipulationModule + from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor + + +@dataclass(frozen=True) +class OperatorStatus: + """Compact dynamic manipulation status.""" + + state: str + error: str + has_plan: bool + + +@dataclass(frozen=True) +class JointTargetRequest: + """Canonical selected joint target request.""" + + group_ids: tuple[PlanningGroupID, ...] + target: JointState + + +@dataclass(frozen=True) +class PoseTargetRequest: + """Explicit world-frame pose target request.""" + + pose_targets: Mapping[PlanningGroupID, PoseStamped] + auxiliary_group_ids: tuple[PlanningGroupID, ...] = () + seed: JointState | None = None + + +@dataclass(frozen=True) +class TargetEvaluationResult: + """Advisory selected-domain target evaluation.""" + + success: bool + status: str + message: str + collision_free: bool = False + group_ids: tuple[PlanningGroupID, ...] = () + target_joints: JointState | None = None + group_diagnostics: Mapping[PlanningGroupID, str] = field(default_factory=dict) + group_poses: Mapping[PlanningGroupID, PoseStamped | None] = field(default_factory=dict) + + +class ManipulationOperator: + """Concrete synchronous facade over ManipulationModule and WorldMonitor.""" + + def __init__(self, module: ManipulationModule, world_monitor: WorldMonitor) -> None: + self._module = module + self._world_monitor = world_monitor + + def status(self) -> OperatorStatus: + """Return compact dynamic state without topology or joint telemetry.""" + return OperatorStatus( + state=self._module.get_state(), + error=self._module.get_error(), + has_plan=self._module.has_planned_path(), + ) + + def get_init_joints(self, robot_name: RobotName) -> JointState | None: + """Return the operator-authoritative init joint state for a robot.""" + init = self._module.get_init_joints(robot_name) + return None if init is None else JointState(init) + + def evaluate_joint_target(self, request: JointTargetRequest) -> TargetEvaluationResult: + """Validate and evaluate a canonical global joint target.""" + groups, validation = self._validate_joint_request(request) + if validation is not None: + return validation + assert groups is not None + complete = self._complete_states(groups, request.target) + if complete is None: + return self._invalid(request.group_ids, "Incomplete robot target state") + return self._evaluate_global_target(groups, JointState(request.target), complete) + + def evaluate_pose_target(self, request: PoseTargetRequest) -> TargetEvaluationResult: + """Validate and evaluate explicit world-frame pose targets.""" + group_ids, validation = self._validate_pose_request(request) + if validation is not None: + return validation + ik = self._module.inverse_kinematics( + pose_targets=dict(request.pose_targets), + auxiliary_group_ids=request.auxiliary_group_ids, + seed=JointState(request.seed) if request.seed is not None else None, + check_collision=True, + ) + if not ik.is_success() or ik.joint_state is None: + return TargetEvaluationResult( + success=False, + status=ik.status.name, + message=ik.message, + collision_free=False, + group_ids=group_ids, + ) + groups = self._groups_for_ids(group_ids) + if groups is None: + return self._invalid(group_ids, "Unknown planning group") + return self._evaluate_global_target(groups, ik.joint_state) + + def plan_to_joints(self, request: JointTargetRequest) -> GeneratedPlan | None: + groups, validation = self._validate_joint_request(request) + if validation is not None: + return None + assert groups is not None + targets = { + group.id: JointState( + { + "name": list(group.joint_names), + "position": list( + request.target.position[offset : offset + len(group.joint_names)] + ), + } + ) + for group, offset in self._group_offsets(groups) + } + return self._module.generate_plan_to_joint_targets( + cast("Mapping[PlanningGroupID | PlanningGroup, JointState]", targets) + ) + + def plan_to_pose(self, request: PoseTargetRequest) -> GeneratedPlan | None: + group_ids, validation = self._validate_pose_request(request) + if validation is not None: + return None + poses = {group_id: stamped for group_id, stamped in request.pose_targets.items()} + return self._module.generate_plan_to_pose_targets( + cast("Mapping[PlanningGroupID | PlanningGroup, PoseStamped]", poses), + request.auxiliary_group_ids, + ) + + def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: + return self._module.preview_plan(plan=plan, duration=duration) + + def execute(self, plan: GeneratedPlan) -> bool: + return self._module.execute_plan(plan=plan) + + def cancel(self) -> bool: + return self._module.cancel() + + def clear_plan(self) -> bool: + return self._module.clear_planned_path() + + def reset(self) -> bool: + result = self._module.reset() + return result.is_success() + + def _validate_joint_request( + self, request: JointTargetRequest + ) -> tuple[tuple[PlanningGroup, ...] | None, TargetEvaluationResult | None]: + groups = self._groups_for_ids(request.group_ids) + if groups is None: + return None, self._invalid( + request.group_ids, "Unknown, duplicate, or overlapping planning group" + ) + expected = tuple(name for group in groups for name in group.joint_names) + names = tuple(str(name) for name in request.target.name) + positions = tuple(float(value) for value in request.target.position) + if len(names) != len(positions): + return None, self._invalid(request.group_ids, "Joint target names and positions differ") + if len(set(names)) != len(names): + return None, self._invalid(request.group_ids, "Joint target contains duplicate joints") + if names != expected: + return None, self._invalid( + request.group_ids, "Joint target must use exact selected global joints in order" + ) + if any(not math.isfinite(value) for value in positions): + return None, self._invalid( + request.group_ids, "Joint target contains non-finite positions" + ) + return groups, None + + def _validate_pose_request( + self, request: PoseTargetRequest + ) -> tuple[tuple[PlanningGroupID, ...], TargetEvaluationResult | None]: + if not request.pose_targets: + return (), self._invalid((), "No pose target") + group_ids = tuple( + dict.fromkeys((*request.pose_targets.keys(), *request.auxiliary_group_ids)) + ) + groups = self._groups_for_ids(group_ids) + if groups is None: + return group_ids, self._invalid( + group_ids, "Unknown, duplicate, or overlapping planning group" + ) + pose_group_ids = set(request.pose_targets) + for group in groups: + if group.id in pose_group_ids and not group.has_pose_target: + return group_ids, self._invalid( + group_ids, f"Planning group '{group.id}' has no tip_link" + ) + for group_id, pose in request.pose_targets.items(): + if pose.frame_id != "world": + return group_ids, self._invalid( + group_ids, f"Unsupported pose frame for '{group_id}': {pose.frame_id}" + ) + if not self._pose_is_finite(pose): + return group_ids, self._invalid( + group_ids, f"Malformed pose target for '{group_id}'" + ) + if request.seed is not None: + seed_names = tuple(str(name) for name in request.seed.name) + if len(seed_names) != len(request.seed.position) or len(set(seed_names)) != len( + seed_names + ): + return group_ids, self._invalid(group_ids, "Malformed seed") + expected = tuple(name for group in groups for name in group.joint_names) + if seed_names != expected or any("/" not in name for name in seed_names): + return group_ids, self._invalid( + group_ids, "Seed must use exact selected global joints in order" + ) + if any(not math.isfinite(float(value)) for value in request.seed.position): + return group_ids, self._invalid(group_ids, "Seed contains non-finite positions") + return group_ids, None + + def _groups_for_ids( + self, group_ids: Sequence[PlanningGroupID] + ) -> tuple[PlanningGroup, ...] | None: + if not group_ids or len(set(group_ids)) != len(group_ids): + return None + try: + selection = self._world_monitor.planning_groups.select(tuple(group_ids)) + except (KeyError, ValueError): + return None + return selection.groups + + def _complete_states( + self, groups: Sequence[PlanningGroup], target: JointState + ) -> dict[RobotName, JointState] | None: + values = { + str(name): float(value) + for name, value in zip(target.name, target.position, strict=True) + } + complete: dict[RobotName, JointState] = {} + for robot_name in dict.fromkeys(group.robot_name for group in groups): + config = self._module.get_robot_config(robot_name) + robot_id = self._module.robot_id_for_name(robot_name) + baseline = ( + None if robot_id is None else self._world_monitor.get_current_joint_state(robot_id) + ) + if config is None or baseline is None or len(baseline.name) != len(baseline.position): + return None + baseline_values = { + str(name): float(value) + for name, value in zip(baseline.name, baseline.position, strict=True) + } + positions: list[float] = [] + for local_name in config.joint_names: + global_name = f"{robot_name}/{local_name}" + value = values.get(global_name, baseline_values.get(local_name)) + if value is None: + return None + positions.append(value) + complete[robot_name] = JointState( + {"name": list(config.joint_names), "position": positions} + ) + return complete + + def _evaluate_global_target( + self, + groups: Sequence[PlanningGroup], + target: JointState, + complete_states: Mapping[RobotName, JointState] | None = None, + ) -> TargetEvaluationResult: + complete = complete_states or self._complete_states(groups, target) + if complete is None: + return self._invalid( + tuple(group.id for group in groups), "Incomplete robot target state" + ) + diagnostics: dict[PlanningGroupID, str] = {} + poses: dict[PlanningGroupID, PoseStamped | None] = {} + valid = True + for group in groups: + robot_id = self._module.robot_id_for_name(group.robot_name) + state = complete[group.robot_name] + group_valid = bool( + robot_id is not None and self._world_monitor.is_state_valid(robot_id, state) + ) + valid = valid and group_valid + diagnostics[group.id] = ( + "Target is collision-free for this robot" + if group_valid + else "Target is in collision or violates limits" + ) + try: + poses[group.id] = self._world_monitor.get_group_ee_pose(group.id, state) + except ValueError: + poses[group.id] = None + return TargetEvaluationResult( + success=valid, + status="FEASIBLE" if valid else "COLLISION", + message="Target is collision-free for each robot" if valid else "Target is infeasible", + collision_free=valid, + group_ids=tuple(group.id for group in groups), + target_joints=JointState(target), + group_diagnostics=diagnostics, + group_poses=poses, + ) + + @staticmethod + def _invalid(group_ids: Sequence[PlanningGroupID], message: str) -> TargetEvaluationResult: + return TargetEvaluationResult(False, "INVALID", message, group_ids=tuple(group_ids)) + + @staticmethod + def _group_offsets(groups: Sequence[PlanningGroup]) -> tuple[tuple[PlanningGroup, int], ...]: + offsets: list[tuple[PlanningGroup, int]] = [] + offset = 0 + for group in groups: + offsets.append((group, offset)) + offset += len(group.joint_names) + return tuple(offsets) + + @staticmethod + def _pose_is_finite(pose: PoseStamped) -> bool: + values = [*pose.position, *pose.orientation] + return len(values) == 7 and all(math.isfinite(float(value)) for value in values) diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index 0ad698748d..d4cf8db0d6 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -27,9 +27,9 @@ from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( - JointPath, Obstacle, - PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, WorldRobotID, ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec @@ -41,25 +41,25 @@ from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class FakeVisualization: - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: return None def get_visualization_url(self) -> str | None: return None - def publish_visualization(self, ctx: object | None = None) -> None: + def update_state(self, frame: VisualizationStateFrame) -> None: return None - def show_preview(self, robot_id: WorldRobotID) -> None: + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: return None - def hide_preview(self, robot_id: WorldRobotID) -> None: - return None - - def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0) -> None: + def cancel_preview_animation(self) -> None: return None def close(self) -> None: @@ -157,6 +157,12 @@ def get_link_pose( def get_jacobian(self, ctx: object, robot_id: WorldRobotID) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) + def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: + return PoseStamped() + + def get_group_jacobian(self, ctx: object, group_id: str) -> NDArray[np.float64]: + return np.zeros((6, 0), dtype=np.float64) + class FakeVisualizationWorld(FakeWorld, FakeVisualization): pass @@ -171,16 +177,18 @@ def test_config_defaults_to_no_visualization() -> None: def test_config_rejects_unknown_visualization_backend() -> None: with pytest.raises(ValidationError, match="visualization"): - ManipulationModuleConfig(visualization={"backend": "bad"}) + ManipulationModuleConfig.model_validate({"visualization": {"backend": "bad"}}) def test_config_validates_viser_visualization() -> None: - config = ManipulationModuleConfig( - visualization={ - "backend": "viser", - "visualization_host": "0.0.0.0", - "visualization_port": "8096", - "viser_panel_enabled": "false", + config = ManipulationModuleConfig.model_validate( + { + "visualization": { + "backend": "viser", + "visualization_host": "0.0.0.0", + "visualization_port": "8096", + "viser_panel_enabled": "false", + } }, ) @@ -191,7 +199,7 @@ def test_config_validates_viser_visualization() -> None: def test_config_meshcat_requires_world_visualization() -> None: - config = ManipulationModuleConfig(visualization={"backend": "meshcat"}) + config = ManipulationModuleConfig.model_validate({"visualization": {"backend": "meshcat"}}) assert isinstance(config.visualization, MeshcatVisualizationConfig) assert config.visualization.requires_world_visualization is True @@ -235,3 +243,18 @@ def test_create_visualization_meshcat_rejects_non_visualization_world() -> None: world_monitor=world_monitor, manipulation_module=MagicMock(), ) + + +def test_create_viser_visualization_has_group_preview_protocol_without_legacy_path_api() -> None: + pytest.importorskip("viser") + + visualization = create_manipulation_visualization( + ViserVisualizationConfig(), + world=FakeWorld(), + world_monitor=MagicMock(), + manipulation_module=MagicMock(), + ) + + assert isinstance(visualization, VisualizationSpec) + assert isinstance(FakeVisualization(), VisualizationSpec) + assert not hasattr(visualization, "animate_path") diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py new file mode 100644 index 0000000000..2c15ef9326 --- /dev/null +++ b/dimos/manipulation/visualization/test_operator.py @@ -0,0 +1,390 @@ +# 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 tests for the manipulation visualization operator facade.""" + +from pathlib import Path + +from dimos.agents.skill_result import SkillResult +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus +from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, + IKResult, + PlanningGroupID, + RobotName, +) +from dimos.manipulation.visualization.operator import ( + JointTargetRequest, + ManipulationOperator, + PoseTargetRequest, +) +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +def _robot_config( + name: str = "arm", + joint_names: list[str] | None = None, + groups: tuple[PlanningGroup, ...] | None = None, +) -> RobotModelConfig: + joints = joint_names or ["j0", "j1"] + definitions = [ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(joints), + base_link="base", + tip_link="tool", + ) + ] + if groups is not None: + definitions = [ + PlanningGroupDefinition( + name=group.group_name, + joint_names=group.local_joint_names, + base_link=group.base_link, + tip_link=group.tip_link, + ) + for group in groups + ] + return RobotModelConfig( + name=name, + model_path=Path("/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=joints, + base_link="base", + planning_groups=definitions, + ) + + +class FakeModule: + def __init__(self) -> None: + self.state = "COMPLETED" + self.error = "" + self.has_plan = True + self.plan = GeneratedPlan( + group_ids=("arm/manipulator",), + trajectory=JointTrajectory( + joint_names=["arm/j0", "arm/j1"], + points=[TrajectoryPoint(0.0, [0.0, 0.0]), TrajectoryPoint(1.25, [0.4, 0.5])], + ), + path=[JointState({"name": ["arm/j0", "arm/j1"], "position": [0.0, 0.0]})], + status=PlanningStatus.SUCCESS, + ) + self.robot_configs: dict[RobotName, RobotModelConfig] = {"arm": _robot_config()} + self.robot_ids: dict[RobotName, str] = {"arm": "arm_id"} + self.plan_joint_targets: list[dict[PlanningGroupID, JointState]] = [] + self.plan_pose_targets: list[ + tuple[dict[PlanningGroupID, PoseStamped], tuple[PlanningGroupID, ...]] + ] = [] + self.ik_calls: list[ + tuple[ + dict[PlanningGroupID, PoseStamped], tuple[PlanningGroupID, ...], JointState | None + ] + ] = [] + self.plan_success = True + self.preview_success = True + self.execute_success = True + self.cancel_success = True + self.clear_success = True + self.reset_success = True + self.topology_calls = 0 + self.telemetry_calls = 0 + + def get_state(self) -> str: + return self.state + + def get_error(self) -> str: + return self.error + + def has_planned_path(self) -> bool: + return self.has_plan + + def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: + self.topology_calls += 1 + return self.robot_configs.get(robot_name) + + def robot_id_for_name(self, robot_name: RobotName) -> str | None: + return self.robot_ids.get(robot_name) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + self.telemetry_calls += 1 + return JointState(name=[f"{robot_name}/j0"], position=[0.0]) + + def inverse_kinematics( + self, + pose_targets: dict[PlanningGroupID, PoseStamped], + auxiliary_group_ids: tuple[PlanningGroupID, ...] = (), + seed: JointState | None = None, + check_collision: bool = True, + ) -> IKResult: + assert check_collision is True + self.ik_calls.append((pose_targets, auxiliary_group_ids, seed)) + return IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState(name=["arm/j0", "arm/j1"], position=[0.4, 0.5]), + message="ok", + ) + + def plan_to_joint_targets(self, targets: dict[PlanningGroupID, JointState]) -> bool: + self.plan_joint_targets.append(targets) + return self.plan_success + + def generate_plan_to_joint_targets( + self, targets: dict[PlanningGroupID, JointState] + ) -> GeneratedPlan | None: + self.plan_joint_targets.append(targets) + return self.plan if self.plan_success else None + + def plan_to_pose_targets( + self, + targets: dict[PlanningGroupID, PoseStamped], + auxiliary_groups: tuple[PlanningGroupID, ...] = (), + ) -> bool: + self.plan_pose_targets.append((targets, auxiliary_groups)) + return self.plan_success + + def generate_plan_to_pose_targets( + self, + targets: dict[PlanningGroupID, PoseStamped], + auxiliary_groups: tuple[PlanningGroupID, ...] = (), + ) -> GeneratedPlan | None: + self.plan_pose_targets.append((targets, auxiliary_groups)) + return self.plan if self.plan_success else None + + def preview_plan( + self, plan: GeneratedPlan | None = None, duration: float | None = None + ) -> bool: + _ = plan + _ = duration + return self.preview_success + + def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: + _ = plan + return self.execute_success + + def cancel(self) -> bool: + return self.cancel_success + + def clear_planned_path(self) -> bool: + return self.clear_success + + def reset(self) -> SkillResult[str]: + if self.reset_success: + return SkillResult.ok("reset") + return SkillResult.fail("ERR", "no reset") + + +class FakeWorldMonitor: + def __init__(self, registry: PlanningGroupRegistry) -> None: + self.planning_groups = registry + self.current_states: dict[str, JointState] = { + "arm_id": JointState(name=["j0", "j1"], position=[0.0, 0.0]) + } + self.valid = True + self.cancel_preview_calls = 0 + self.telemetry_calls = 0 + + def get_current_joint_state(self, robot_id: str) -> JointState | None: + self.telemetry_calls += 1 + return self.current_states.get(robot_id) + + def is_state_valid(self, robot_id: str, joint_state: JointState) -> bool: + _ = robot_id, joint_state + return self.valid + + def get_group_ee_pose( + self, group_id: PlanningGroupID, joint_state: JointState | None = None + ) -> PoseStamped: + _ = group_id, joint_state + return PoseStamped( + frame_id="world", position=Vector3(1.0, 2.0, 3.0), orientation=Quaternion() + ) + + def cancel_preview_animation(self) -> None: + self.cancel_preview_calls += 1 + + +def _operator( + config: RobotModelConfig | None = None, +) -> tuple[ManipulationOperator, FakeModule, FakeWorldMonitor]: + robot_config = config or _robot_config() + module = FakeModule() + module.robot_configs = {robot_config.name: robot_config} + module.robot_ids = {robot_config.name: f"{robot_config.name}_id"} + monitor = FakeWorldMonitor(PlanningGroupRegistry([robot_config])) + monitor.current_states = { + f"{robot_config.name}_id": JointState( + name=robot_config.joint_names, position=[0.0] * len(robot_config.joint_names) + ) + } + return ManipulationOperator(module, monitor), module, monitor # type: ignore[arg-type] + + +def _joint_request( + names: list[str] | None = None, positions: list[float] | None = None +) -> JointTargetRequest: + return JointTargetRequest( + group_ids=("arm/manipulator",), + target=JointState(name=names or ["arm/j0", "arm/j1"], position=positions or [0.1, 0.2]), + ) + + +def _pose(frame_id: str = "world") -> PoseStamped: + return PoseStamped(frame_id=frame_id, position=Vector3(0.1, 0.2, 0.3), orientation=Quaternion()) + + +def test_status_is_compact_and_does_not_read_topology_or_telemetry() -> None: + operator, module, monitor = _operator() + + status = operator.status() + + assert status.state == "COMPLETED" + assert status.error == "" + assert status.has_plan is True + assert module.topology_calls == 0 + assert module.telemetry_calls == 0 + assert monitor.telemetry_calls == 0 + + +def test_evaluate_joint_target_accepts_exact_global_selection_domain() -> None: + operator, _, _ = _operator() + + result = operator.evaluate_joint_target(_joint_request()) + + assert result.success is True + assert result.status == "FEASIBLE" + assert result.target_joints is not None + assert list(result.target_joints.name) == ["arm/j0", "arm/j1"] + assert list(result.target_joints.position) == [0.1, 0.2] + assert result.group_diagnostics["arm/manipulator"] == "Target is collision-free for this robot" + assert result.group_poses["arm/manipulator"] is not None + + +def test_joint_target_validation_rejects_bad_joint_requests() -> None: + cases = [ + _joint_request(["j0", "j1"], [0.1, 0.2]), + _joint_request(["arm/j0", "arm/j0"], [0.1, 0.2]), + _joint_request(["arm/j0"], [0.1]), + _joint_request(["arm/j0", "arm/j1", "arm/extra"], [0.1, 0.2, 0.3]), + _joint_request(["arm/j0", "arm/j1"], [0.1, float("nan")]), + JointTargetRequest( + ("missing/manipulator",), JointState(name=["missing/j0"], position=[0.1]) + ), + JointTargetRequest( + ("arm/manipulator", "arm/manipulator"), + JointState(name=["arm/j0", "arm/j1"], position=[0.1, 0.2]), + ), + ] + operator, _, _ = _operator() + + for request in cases: + result = operator.evaluate_joint_target(request) + assert result.success is False + assert result.status == "INVALID" + + +def test_joint_target_validation_rejects_overlapping_groups() -> None: + groups = ( + PlanningGroup("arm/first", "arm", "first", ("arm/j0",), ("j0",), "base"), + PlanningGroup("arm/second", "arm", "second", ("arm/j0",), ("j0",), "base"), + ) + operator, _, _ = _operator(_robot_config(groups=groups)) + request = JointTargetRequest( + ("arm/first", "arm/second"), JointState(name=["arm/j0", "arm/j0"], position=[0.1, 0.2]) + ) + + result = operator.evaluate_joint_target(request) + + assert result.success is False + assert result.status == "INVALID" + + +def test_pose_evaluation_accepts_world_frame_and_delegates_original_request() -> None: + operator, module, _ = _operator() + pose = _pose() + seed = JointState(name=["arm/j0", "arm/j1"], position=[0.0, 0.0]) + request = PoseTargetRequest({"arm/manipulator": pose}, seed=seed) + + result = operator.evaluate_pose_target(request) + + assert result.success is True + assert result.target_joints is not None + assert list(result.target_joints.name) == ["arm/j0", "arm/j1"] + assert module.ik_calls == [({"arm/manipulator": pose}, (), seed)] + + +def test_pose_validation_rejects_frame_capability_and_seed_errors() -> None: + no_pose_group = ( + PlanningGroup("arm/no_pose", "arm", "no_pose", ("arm/j0",), ("j0",), "base", None), + ) + no_pose_operator, _, _ = _operator(_robot_config(joint_names=["j0"], groups=no_pose_group)) + bad_seed_cases = [ + PoseTargetRequest({"arm/manipulator": _pose("camera")}), + PoseTargetRequest( + {"arm/manipulator": _pose()}, + seed=JointState(name=["j0", "j1"], position=[0.0, 0.0]), + ), + PoseTargetRequest( + {"arm/manipulator": _pose()}, + seed=JointState(name=["arm/j0", "arm/j0"], position=[0.0, 0.0]), + ), + ] + operator, _, _ = _operator() + + no_pose = no_pose_operator.evaluate_pose_target(PoseTargetRequest({"arm/no_pose": _pose()})) + assert no_pose.success is False + assert no_pose.status == "INVALID" + for request in bad_seed_cases: + result = operator.evaluate_pose_target(request) + assert result.success is False + assert result.status == "INVALID" + + +def test_planning_methods_return_exact_generated_plan() -> None: + operator, module, _ = _operator() + joint_request = _joint_request() + pose = _pose() + pose_request = PoseTargetRequest({"arm/manipulator": pose}) + + joint_result = operator.plan_to_joints(joint_request) + pose_result = operator.plan_to_pose(pose_request) + + assert joint_result is module.plan + assert list(module.plan_joint_targets[0]["arm/manipulator"].name) == ["arm/j0", "arm/j1"] + assert module.plan_pose_targets == [({"arm/manipulator": pose}, ())] + assert pose_result is module.plan + + +def test_actions_return_typed_results_and_cancel_fallback_ownership() -> None: + operator, module, monitor = _operator() + + assert operator.preview(module.plan, 0.5) is True + assert operator.execute(module.plan) is True + assert operator.clear_plan() is True + assert operator.reset() is True + cancel_result = operator.cancel() + assert cancel_result is True + assert monitor.cancel_preview_calls == 0 + + module.cancel_success = False + fallback = operator.cancel() + assert fallback is False + assert monitor.cancel_preview_calls == 0 diff --git a/dimos/manipulation/visualization/viser/adapter.py b/dimos/manipulation/visualization/viser/adapter.py deleted file mode 100644 index a3f363d8e4..0000000000 --- a/dimos/manipulation/visualization/viser/adapter.py +++ /dev/null @@ -1,174 +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. - -from __future__ import annotations - -from collections.abc import Sequence -from typing import TYPE_CHECKING - -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.msgs.sensor_msgs.JointState import JointState - -if TYPE_CHECKING: - from dimos.manipulation.manipulation_module import ManipulationModule - from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor - from dimos.manipulation.planning.spec.config import RobotModelConfig - from dimos.manipulation.planning.spec.models import RobotName, WorldRobotID - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - - -def copy_joint_state(joint_state: JointState | None) -> JointState | None: - """Make a local copy of a JointState-like message for rendering.""" - return None if joint_state is None else JointState(joint_state) - - -def _string_list(value: object) -> list[str]: - return [str(item) for item in value] if isinstance(value, list) else [] - - -def _float_list_or_none(value: object) -> list[float] | None: - return [float(item) for item in value] if isinstance(value, list) else None - - -def _float_value(value: object, default: float = 0.0) -> float: - return float(value) if isinstance(value, int | float | str) else default - - -class InProcessViserAdapter: - """Small in-process boundary between Viser callbacks and manipulation internals.""" - - def __init__( - self, - *, - world_monitor: WorldMonitor, - manipulation_module: ManipulationModule, - ) -> None: - self._world_monitor = world_monitor - self._module = manipulation_module - - def list_robots(self) -> list[RobotName]: - return list(self._module.list_robots()) - - def robot_items(self) -> list[tuple[RobotName, WorldRobotID, RobotModelConfig]]: - return self._module.robot_items() - - def robot_id_for_name(self, robot_name: RobotName) -> WorldRobotID | None: - return self._module.robot_id_for_name(robot_name) - - def robot_name_for_id(self, robot_id: WorldRobotID) -> RobotName | None: - return self._module.robot_name_for_id(robot_id) - - def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: - return self._module.get_robot_config(robot_name) - - def get_robot_info(self, robot_name: RobotName) -> RobotInfo | None: - info = self._module.get_robot_info(robot_name) - if info is None: - return None - - robot_info: RobotInfo = { - "name": str(info["name"]), - "world_robot_id": str(info["world_robot_id"]), - "joint_names": _string_list(info["joint_names"]), - "end_effector_link": str(info["end_effector_link"]), - "base_link": str(info["base_link"]), - "max_velocity": _float_value(info["max_velocity"]), - "max_acceleration": _float_value(info["max_acceleration"]), - "has_joint_name_mapping": bool(info["has_joint_name_mapping"]), - "coordinator_task_name": None - if info["coordinator_task_name"] is None - else str(info["coordinator_task_name"]), - "home_joints": _float_list_or_none(info["home_joints"]), - "pre_grasp_offset": _float_value(info["pre_grasp_offset"]), - "init_joints": _float_list_or_none(info["init_joints"]), - } - return robot_info - - def get_init_joints(self, robot_name: RobotName) -> JointState | None: - return copy_joint_state(self._module.get_init_joints(robot_name)) - - def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None: - return None - return copy_joint_state(self._world_monitor.get_current_joint_state(robot_id)) - - def is_state_stale(self, robot_name: RobotName, max_age: float = 1.0) -> bool: - robot_id = self.robot_id_for_name(robot_name) - return True if robot_id is None else self._world_monitor.is_state_stale(robot_id, max_age) - - def get_ee_pose( - self, robot_name: RobotName, joint_state: JointState | None = None - ) -> PoseStamped | None: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None: - return None - return self._world_monitor.get_ee_pose(robot_id, copy_joint_state(joint_state)) - - def evaluate_joint_target(self, joints: JointState, robot_name: RobotName) -> TargetEvaluation: - """Evaluate a joint target through WorldMonitor helpers, not raw WorldSpec access.""" - result: TargetEvaluation = { - **self._module.evaluate_joint_target(copy_joint_state(joints), robot_name) - } - joint_state = result.get("joint_state") - result["joint_state"] = copy_joint_state( - joint_state if isinstance(joint_state, JointState) else None - ) - return result - - def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvaluation: - """Evaluate a Cartesian target through module/WorldMonitor helper boundaries.""" - result: TargetEvaluation = {**self._module.evaluate_pose_target(pose, robot_name)} - joint_state = result.get("joint_state") - result["joint_state"] = copy_joint_state( - joint_state if isinstance(joint_state, JointState) else None - ) - return result - - def get_module_state(self) -> str: - return str(self._module.get_state()) - - def get_error(self) -> str: - return self._module.get_error() - - def reset(self) -> bool: - return self._module.reset().is_success() - - def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: - return self._module.plan_to_pose(pose, robot_name) - - def plan_to_joints(self, joints: JointState, robot_name: RobotName | None = None) -> bool: - return self._module.plan_to_joints(joints, robot_name) - - def preview_path(self, robot_name: RobotName | None = None) -> bool: - return self._module.preview_path(robot_name=robot_name) - - def execute(self, robot_name: RobotName | None = None) -> bool: - return self._module.execute(robot_name) - - def cancel(self) -> bool: - return self._module.cancel() - - def clear_planned_path(self) -> bool: - return self._module.clear_planned_path() - - @staticmethod - def joints_from_values(joint_names: Sequence[str], values: Sequence[float]) -> JointState: - return JointState( - { - "name": list(joint_names), - "position": [float(value) for value in values], - } - ) diff --git a/dimos/manipulation/visualization/viser/animation.py b/dimos/manipulation/visualization/viser/animation.py index f5a574d82c..6754343e21 100644 --- a/dimos/manipulation/visualization/viser/animation.py +++ b/dimos/manipulation/visualization/viser/animation.py @@ -14,85 +14,49 @@ from __future__ import annotations -from collections.abc import Callable, Sequence -import time +from collections.abc import Sequence +from dataclasses import dataclass +from itertools import pairwise -from dimos.msgs.sensor_msgs.JointState import JointState +@dataclass(frozen=True) +class PreviewFrame: + """One timestamped local robot preview frame.""" -def interpolate_joint_path( - path: Sequence[JointState], duration: float, fps: float -) -> list[list[float]]: - """Interpolate a joint path into visualization frames.""" - waypoints = [list(waypoint.position) for waypoint in path if waypoint.position] - if not waypoints: - return [] - if len(waypoints) == 1 or duration <= 0.0: - return [waypoints[-1]] - frame_count = max(int(duration * max(fps, 1.0)) + 1, len(waypoints)) - segment_count = len(waypoints) - 1 - frames: list[list[float]] = [] - for frame_index in range(frame_count): - path_t = frame_index / max(frame_count - 1, 1) - scaled = path_t * segment_count - segment_index = min(int(scaled), segment_count - 1) - local_t = scaled - segment_index - start = waypoints[segment_index] - end = waypoints[segment_index + 1] - if len(start) != len(end): - continue - frames.append( - [ - start_value + (end_value - start_value) * local_t - for start_value, end_value in zip(start, end, strict=False) - ] - ) - if frames and frames[-1] != waypoints[-1]: - frames.append(waypoints[-1]) - return frames + time_from_start: float + positions: tuple[float, ...] -def sampled_joint_path_frames( - path: Sequence[JointState], duration: float, fps: float -) -> list[list[float]]: - """Return animation frames while preserving already sampled trajectories. +@dataclass(frozen=True) +class PreviewTrack: + """One fixed-baseline local robot track in a group-native preview.""" - ManipulationModule.preview_path() owns trajectory-aware interpolation because it has access - to JointTrajectory waypoint timing. If a path arrives already sampled near the target display - rate, Viser should play those samples directly instead of re-interpolating by waypoint index. - Sparse direct VisualizationSpec callers still get local interpolation as a fallback. - """ - waypoints = [list(waypoint.position) for waypoint in path if waypoint.position] - if not waypoints: - return [] - expected_frames = max(int(duration * max(fps, 1.0)) + 1, 1) if duration > 0.0 else 1 - if len(waypoints) >= expected_frames: - return waypoints - return interpolate_joint_path(path, duration, fps) + robot_id: str + joint_names: tuple[str, ...] + frames: tuple[PreviewFrame, ...] -class PreviewAnimator: - """Blocking preview-ghost path animator with Meshcat-compatible semantics. +@dataclass(frozen=True) +class GroupPreviewAnimation: + """Validated collection of robot tracks sharing one preview transaction.""" - This class is only for transient path playback. Persistent target ghosts are updated - directly by scene target methods and must not be routed through this animator. - """ + tracks: tuple[PreviewTrack, ...] - def __init__( - self, - set_joints: Callable[[Sequence[float]], None], - *, - sleep: Callable[[float], None] = time.sleep, - ) -> None: - self._set_joints = set_joints - self._sleep = sleep - def animate(self, path: Sequence[JointState], duration: float, fps: float) -> bool: - frames = sampled_joint_path_frames(path, duration, fps) - if not frames: - return False - step_delay = duration / max(len(frames) - 1, 1) if duration > 0.0 else 0.0 - for joints in frames: - self._set_joints(joints) - self._sleep(step_delay) - return True +def scaled_frame_delays(frames: Sequence[PreviewFrame], duration: float) -> tuple[float, ...]: + """Return stored inter-frame delays, optionally scaled to a requested duration.""" + if len(frames) < 2: + return () + original_duration = max(float(frames[-1].time_from_start), 0.0) + scale = duration / original_duration if duration > 0.0 and original_duration > 0.0 else 1.0 + return tuple( + max(float(next_frame.time_from_start) - float(frame.time_from_start), 0.0) * scale + for frame, next_frame in pairwise(frames) + ) + + +def preview_tick_times(preview: GroupPreviewAnimation) -> tuple[float, ...]: + """Union all stored track timestamps without synthesizing extra samples.""" + return tuple( + sorted({float(frame.time_from_start) for track in preview.tracks for frame in track.frames}) + ) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index 218290c5c4..56eb548ee9 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -14,10 +14,18 @@ from __future__ import annotations -from typing import TypeAlias - -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from collections.abc import Mapping, MutableMapping, Sequence +from typing import TypeAlias, cast + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import PlanningGroupID, PlanningSceneInfo, RobotName +from dimos.manipulation.visualization.operator import ( + JointTargetRequest, + ManipulationOperator, + PoseTargetRequest, + TargetEvaluationResult, +) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.runtime import VISER_INSTALL_HINT from dimos.manipulation.visualization.viser.scene import ViserManipulationScene @@ -35,6 +43,7 @@ TargetStatus, ) from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -68,28 +77,52 @@ # Fallback joint-slider range (radians) when a robot config omits joint limits. DEFAULT_JOINT_LIMITS = (-3.14, 3.14) +PRIMARY_ACTION_COLOR = (0, 102, 179) +ACTIVE_GROUP_COLOR = PRIMARY_ACTION_COLOR +INACTIVE_GROUP_COLOR = (52, 52, 52) + + +def group_display_name(group: PlanningGroup) -> str: + return ( + str(group.robot_name) + if str(group.group_name) == "manipulator" + else f"{group.robot_name} {group.group_name}" + ) + + +def _copy_joint_state(state: JointState | None) -> JointState | None: + return None if state is None else JointState(state) class ViserPanelGui: - """Optional operator panel with parity for the original cc/viser-vis panel.""" + """Viser operator panel for manipulation target editing and plan control.""" def __init__( self, server: ViserServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: ManipulationOperator | object, + current_states: MutableMapping[str, JointState], config: ViserVisualizationConfig, scene: ViserManipulationScene | None = None, ) -> None: self.server = server - self.adapter = adapter + self.scene_info = scene_info + self.operator = cast("ManipulationOperator", operator) + self.current_states = current_states + self._robots_by_name = { + config.name: (robot_id, config) for robot_id, config in scene_info.robots.items() + } + self._scene_groups_by_id = {group.id: group for group in scene_info.planning_groups} self.config = config self.scene = scene self.state = PanelState(runtime=PanelRuntime.STARTING) self._closed = False self._operation_sequence_id = 0 self._suppress_target_callbacks = False + self._default_group_initialized = False self._handles: dict[str, PanelHandle] = {} - self._joint_sliders: dict[str, GuiSliderHandle[float]] = {} + self._joint_sliders: dict[tuple[PlanningGroupID, str], GuiSliderHandle[float]] = {} self._worker = TargetEvaluationWorker( self._handle_target_evaluation_request, self._apply_target_evaluation_result, @@ -117,6 +150,8 @@ def close(self) -> None: return self._closed = True self.state.runtime = PanelRuntime.STOPPING + if self.scene is not None: + self.scene.cancel_preview_animation() self._worker.stop() self._operation_worker.stop(timeout=2.0) self._clear_joint_sliders() @@ -124,20 +159,171 @@ def close(self) -> None: self._handles.clear() self.state.runtime = PanelRuntime.STOPPED + def list_robots(self) -> list[RobotName]: + return [config.name for config in self.scene_info.robots.values()] + + def list_planning_groups(self) -> list[PlanningGroup]: + return list(self.scene_info.planning_groups) + + def robot_items(self) -> list[tuple[RobotName, str, RobotModelConfig]]: + return [ + (config.name, str(robot_id), config) + for robot_id, config in self.scene_info.robots.items() + ] + + def robot_id_for_name(self, robot_name: RobotName) -> str | None: + item = self._robots_by_name.get(robot_name) + return None if item is None else str(item[0]) + + def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: + item = self._robots_by_name.get(robot_name) + return None if item is None else item[1] + + def get_init_joints(self, robot_name: RobotName) -> JointState | None: + init = self.operator.get_init_joints(robot_name) + if init is None: + return None + config = self.get_robot_config(robot_name) + if config is None: + return JointState(init) + values = self._local_values_for_robot(robot_name, init) + if any(name not in values for name in config.joint_names): + return JointState(init) + return JointState( + { + "name": list(config.joint_names), + "position": [values[name] for name in config.joint_names], + } + ) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + robot_id = self.robot_id_for_name(robot_name) + return None if robot_id is None else _copy_joint_state(self.current_states.get(robot_id)) + + def get_group_ee_pose(self, group_id: PlanningGroupID) -> PoseStamped | None: + group = self._scene_groups_by_id.get(group_id) + if group is None: + return None + targets = self._current_target_for_group(group) + if group_id not in targets: + return None + return self.evaluate_joint_target_set((group_id,), targets).group_poses.get(group_id) + + def _current_target_for_group(self, group: PlanningGroup) -> dict[PlanningGroupID, JointState]: + current = self.get_current_joint_state(group.robot_name) + if current is None or len(current.name) != len(current.position): + return {} + values = self._local_values_for_robot(group.robot_name, current) + if any(name not in values for name in group.local_joint_names): + return {} + return { + group.id: JointState( + { + "name": list(group.joint_names), + "position": [values[name] for name in group.local_joint_names], + } + ) + } + + def is_state_stale(self, robot_name: RobotName, max_age: float = 1.0) -> bool: + _ = max_age + return self.get_current_joint_state(robot_name) is None + + def get_module_state(self) -> str: + return self.operator.status().state + + def get_error(self) -> str: + return self.operator.status().error + + def reset(self) -> bool: + return self.operator.reset() + + def evaluate_joint_target_set( + self, group_ids: Sequence[PlanningGroupID], targets: Mapping[PlanningGroupID, JointState] + ) -> TargetEvaluationResult: + names: list[str] = [] + positions: list[float] = [] + for group_id in group_ids: + target = targets.get(group_id) + if target is None: + return TargetEvaluationResult( + False, "INVALID", "Incomplete joint target", group_ids=tuple(group_ids) + ) + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + return self.operator.evaluate_joint_target( + JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + ) + + def evaluate_pose_target_set( + self, + pose_targets: Mapping[PlanningGroupID, Pose], + auxiliary_group_ids: Sequence[PlanningGroupID] = (), + seed: JointState | None = None, + ) -> TargetEvaluationResult: + stamped = { + group_id: PoseStamped( + frame_id="world", position=pose.position, orientation=pose.orientation + ) + for group_id, pose in pose_targets.items() + } + return self.operator.evaluate_pose_target( + PoseTargetRequest(stamped, tuple(auxiliary_group_ids), _copy_joint_state(seed)) + ) + + def cancel(self) -> bool: + return self.operator.cancel() + + def clear_planned_path(self) -> bool: + return self.operator.clear_plan() + + def plan_to_selected_joints( + self, group_ids: Sequence[PlanningGroupID], targets: Mapping[PlanningGroupID, JointState] + ) -> bool: + names: list[str] = [] + positions: list[float] = [] + for group_id in group_ids: + target = targets.get(group_id) + if target is None: + return False + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + plan = self.operator.plan_to_joints( + JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + ) + self.state.plan_state.plan = plan + return plan is not None + + def preview_path(self) -> bool: + plan = self.state.plan_state.plan + return plan is not None and self.operator.preview(plan) + + def execute(self) -> bool: + plan = self.state.plan_state.plan + return plan is not None and self.operator.execute(plan) + def refresh(self) -> None: if self._closed: return - robots = self.adapter.list_robots() + robots = self.list_robots() + groups = self.list_planning_groups() self.state.backend_status = ( BackendConnectionStatus.READY if robots else BackendConnectionStatus.WAITING_FOR_ROBOT ) - if self.state.selected_robot is None and robots: - self.state.selected_robot = robots[0] + if not self.state.selected_group_ids and groups and not self._default_group_initialized: + first = next((group for group in groups if group.has_pose_target), groups[0]) + self.state.selected_group_ids = (first.id,) + self.state.selected_robot = str(first.robot_name) self.state.target_status = TargetStatus.EMPTY + self._default_group_initialized = True + initialized_groups = set(self.state.group_joint_targets) + self._initialize_selected_group_targets() + if set(self.state.group_joint_targets) != initialized_groups: self._build_joint_sliders() - self._sync_robot_dropdown(robots) + self._sync_group_selector(groups) self._refresh_selected_robot_state() self._ensure_scene_controls() + self._sync_target_ghost_visibility() self._sync_preset_dropdown() self._update_status_text() self._update_control_state() @@ -150,26 +336,26 @@ def _build(self) -> None: self._build_panel_controls(gui) def _build_panel_controls(self, gui: GuiApi) -> None: - self._handles["status"] = gui.add_markdown("Starting manipulation panel...") - robots = self.adapter.list_robots() + self._handles["status"] = gui.add_markdown("### Status\n**State:** Ready") self._build_scene_controls(gui) - robot_dropdown = gui.add_dropdown( - "Robot", - options=robots or [""], - initial_value=robots[0] if robots else "", + self._handles["planning_groups_heading"] = gui.add_markdown( + "### Planning Groups\nActive planning groups for pose goals, planning, and joint edits." ) - robot_dropdown.on_update(lambda event: self._select_robot(event.target.value)) - self._handles["robot"] = robot_dropdown + self._sync_group_selector(self.list_planning_groups()) + self._handles["target_heading"] = gui.add_markdown("### Target") preset_dropdown = gui.add_dropdown( - "Target Preset", + "Preset", options=["Select preset...", "Current"], initial_value="Select preset...", ) preset_dropdown.on_update(lambda event: self._apply_preset(event.target.value)) self._handles["preset"] = preset_dropdown - plan_button = gui.add_button("Plan", disabled=True) + self._handles["target_summary"] = gui.add_markdown("Feasibility: `unknown`") + self._handles["actions_heading"] = gui.add_markdown("### Actions") + plan_button = gui.add_button("Plan", disabled=True, color=PRIMARY_ACTION_COLOR) plan_button.on_click(lambda _: self._submit_plan()) self._handles["plan"] = plan_button + self._handles["plan_controls_heading"] = gui.add_markdown("**Plan controls**") preview_button = gui.add_button("Preview", disabled=True) preview_button.on_click(lambda _: self._submit_preview()) self._handles["preview"] = preview_button @@ -182,8 +368,63 @@ def _build_panel_controls(self, gui: GuiApi) -> None: clear_button = gui.add_button("Clear plan") clear_button.on_click(lambda _: self._submit_clear()) self._handles["clear"] = clear_button + joint_controls = gui.add_folder("Joint Control", expand_by_default=False) + self._handles["joint_control_folder"] = joint_controls self._build_joint_sliders() + def _sync_group_selector(self, groups: list[PlanningGroup]) -> None: + """Render source-order group toggle buttons without a robot dropdown.""" + selected = set(self.state.selected_group_ids) + seen: set[str] = set() + for group in sorted( + groups, key=lambda item: (not bool(item.has_pose_target), str(item.id)) + ): + group_id = str(group.id) + key = f"group:{group_id}" + seen.add(key) + label = group_display_name(group) + handle = self._handles.get(key) + color = ACTIVE_GROUP_COLOR if group_id in selected else INACTIVE_GROUP_COLOR + if handle is None: + handle = self.server.gui.add_button( + label, + color=color, + hint="Click to toggle this planning group in the target set.", + ) + + def on_click(_event: object, selected_group_id: str = group_id) -> None: + self._toggle_group_selected(selected_group_id) + + handle.on_click(on_click) + self._handles[key] = handle + else: + self._set_optional_handle_attr(handle, "label", label) + self._set_optional_handle_attr(handle, "color", color) + for key in [key for key in self._handles if key.startswith("group:") and key not in seen]: + handle = self._handles.pop(key) + remove = getattr(handle, "remove", None) + if callable(remove): + remove() + + def _toggle_group_selected(self, group_id: str) -> None: + groups = {str(group.id): group for group in self.list_planning_groups()} + if group_id not in groups: + return + current = list(self.state.selected_group_ids) + if group_id in current: + current.remove(group_id) + else: + current.append(group_id) + self.state.selected_group_ids = tuple(current) + self.state.advance_selection_epoch() + self._clear_invalidated_preview() + first = groups.get(current[0]) if current else None + self.state.selected_robot = None if first is None else str(first.robot_name) + self._prune_inactive_group_state() + self._initialize_selected_group_targets() + self._build_joint_sliders() + self.refresh() + def _build_scene_controls(self, gui: GuiApi) -> None: if self.scene is None: return @@ -203,71 +444,102 @@ def _set_scene_grid_visible(self, visible: bool) -> None: def _refresh_selected_robot_state(self) -> None: robot_name = self.state.selected_robot if robot_name is None: - self.state.robot_info = None self.state.current_joints = None - self.state.current_ee_pose = None - self.state.manipulation_state = self.adapter.get_module_state() + self.state.manipulation_state = self.get_module_state() return - self.state.robot_info = self.adapter.get_robot_info(robot_name) - current = self.adapter.get_current_joint_state(robot_name) + current = self.get_current_joint_state(robot_name) self.state.current_joints = list(current.position) if current is not None else None - self.state.current_ee_pose = self.adapter.get_ee_pose(robot_name) - self.state.manipulation_state = self.adapter.get_module_state() - adapter_error = self.adapter.get_error() + self.state.manipulation_state = self.get_module_state() + adapter_error = self.get_error() if adapter_error: self.state.error = adapter_error def _ensure_scene_controls(self) -> None: - if self.scene is None or self.state.selected_robot is None: - return - robot_id = self.adapter.robot_id_for_name(self.state.selected_robot) - if robot_id is None: + if self.scene is None: return - ee_control = self.scene.ensure_target_controls(str(robot_id), self._on_transform_update) - if ee_control is not None: - self._handles["ee_control"] = ee_control - if ( - self.state.target_status == TargetStatus.EMPTY - and self.state.current_ee_pose is not None - ): - self.state.cartesian_target = self.state.current_ee_pose - self._suppress_target_callbacks = True - try: - self.scene.set_target_pose(str(robot_id), self.state.current_ee_pose) - finally: - self._suppress_target_callbacks = False + groups = self._groups_by_id() + pose_group_ids = tuple( + group_id + for group_id in self.state.selected_group_ids + if (group := groups.get(group_id)) is not None and group.has_pose_target + ) + for key in [key for key in self._handles if key.startswith("ee_control:")]: + if key.removeprefix("ee_control:") not in pose_group_ids: + self.scene.remove_target_controls(key.removeprefix("ee_control:")) + self._handles.pop(key, None) + for group_id in pose_group_ids: + group = groups[group_id] + + def on_transform_update( + target: TransformControlsHandle, + selected_group_id: PlanningGroupID = group_id, + ) -> None: + self._on_transform_update(selected_group_id, target) + + control = self.scene.ensure_target_controls( + str(group_id), + on_transform_update, + ) + if control is not None: + self._handles[f"ee_control:{group_id}"] = control + pose = self.state.pose_targets.get(group_id) + if pose is not None: + self._suppress_target_callbacks = True + try: + self.scene.set_target_pose(str(group_id), pose) + finally: + self._suppress_target_callbacks = False def _build_joint_sliders(self) -> None: - if self.state.selected_robot is None: - return gui = self.server.gui - config = self.adapter.get_robot_config(self.state.selected_robot) - if config is None: - return - current = self.adapter.get_current_joint_state(self.state.selected_robot) - values = list(current.position) if current is not None else [0.0] * len(config.joint_names) self._clear_joint_sliders() - joint_limits_lower = config.joint_limits_lower - joint_limits_upper = config.joint_limits_upper - for index, joint_name in enumerate(config.joint_names): - lower, upper = DEFAULT_JOINT_LIMITS - if joint_limits_lower is not None and index < len(joint_limits_lower): - lower = joint_limits_lower[index] - if joint_limits_upper is not None and index < len(joint_limits_upper): - upper = joint_limits_upper[index] - handle = gui.add_slider( - joint_name, - min=float(lower), - max=float(upper), - step=0.001, - initial_value=float(values[index] if index < len(values) else 0.0), - ) + if not self.state.selected_group_ids: + return + joint_folder = self._handles.get("joint_control_folder") + if joint_folder is not None: + folder = cast("GuiFolderHandle", joint_folder) + with folder: + self._build_joint_slider_handles(gui) + return + self._build_joint_slider_handles(gui) + + def _build_joint_slider_handles(self, gui: GuiApi) -> None: + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + continue + config = self.get_robot_config(group.robot_name) + target = self.state.group_joint_targets.get(group_id) + if config is None or target is None: + continue + config_indexes = {str(name): index for index, name in enumerate(config.joint_names)} + for _global_name, local_name, value in zip( + group.joint_names, group.local_joint_names, target.position, strict=True + ): + index = config_indexes.get(str(local_name)) + lower, upper = DEFAULT_JOINT_LIMITS + if index is not None and config.joint_limits_lower is not None: + lower = config.joint_limits_lower[index] + if index is not None and config.joint_limits_upper is not None: + upper = config.joint_limits_upper[index] + key = (group_id, str(local_name)) + handle = gui.add_slider( + f"{group_id}/{local_name}", + min=float(lower), + max=float(upper), + step=0.001, + initial_value=float(value), + ) - def on_update(_event: object, name: str = joint_name) -> None: - self._on_joint_slider_update(name) + def on_slider_update( + _event: object, + selected_group_id: PlanningGroupID = group_id, + name: str = str(local_name), + ) -> None: + self._on_joint_slider_update(selected_group_id, name) - handle.on_update(on_update) - self._joint_sliders[joint_name] = handle + handle.on_update(on_slider_update) + self._joint_sliders[key] = handle def _clear_joint_sliders(self) -> None: for handle in self._joint_sliders.values(): @@ -277,6 +549,127 @@ def _clear_joint_sliders(self) -> None: pass self._joint_sliders.clear() + def _groups_by_id(self) -> dict[PlanningGroupID, PlanningGroup]: + return {group.id: group for group in self.list_planning_groups()} + + def _selected_robot_names(self) -> tuple[str, ...]: + groups = self._groups_by_id() + return tuple( + dict.fromkeys( + str(groups[group_id].robot_name) + for group_id in self.state.selected_group_ids + if group_id in groups + ) + ) + + def _stale_robot_names(self, group_ids: tuple[PlanningGroupID, ...]) -> tuple[str, ...]: + """Return every affected robot whose monitored joint state is stale.""" + groups = self._groups_by_id() + robot_names = tuple( + dict.fromkeys( + str(groups[group_id].robot_name) for group_id in group_ids if group_id in groups + ) + ) + return tuple(name for name in robot_names if self.is_state_stale(name)) + + def _state_values_by_local_name(self, state: JointState | None) -> dict[str, float]: + if state is None or len(state.name) != len(state.position): + return {} + return { + str(name): float(value) for name, value in zip(state.name, state.position, strict=True) + } + + def _local_values_for_robot( + self, robot_name: str, state: JointState | None + ) -> dict[str, float]: + config = self.get_robot_config(robot_name) + if config is None or state is None or len(state.name) != len(state.position): + return {} + raw = self._state_values_by_local_name(state) + values: dict[str, float] = {} + for local_name in config.joint_names: + global_name = f"{robot_name}/{local_name}" + if local_name in raw: + values[local_name] = raw[local_name] + elif global_name in raw: + values[local_name] = raw[global_name] + return values + + def _initialize_selected_group_targets(self) -> None: + for group_id in self.state.selected_group_ids: + if group_id in self.state.group_joint_targets: + continue + group = self._groups_by_id().get(group_id) + if group is None: + continue + if self.is_state_stale(group.robot_name): + continue + values = self._local_values_for_robot( + str(group.robot_name), self.get_current_joint_state(group.robot_name) + ) + if any(str(name) not in values for name in group.local_joint_names): + continue + self.state.group_joint_targets[group_id] = JointState( + { + "name": list(group.joint_names), + "position": [float(values[str(name)]) for name in group.local_joint_names], + } + ) + if group.has_pose_target and group_id not in self.state.pose_targets: + pose = self.get_group_ee_pose(group_id) + if pose is not None: + self.state.pose_targets[group_id] = pose + self.state.group_poses[group_id] = pose + if self.state.cartesian_target is None: + self.state.cartesian_target = pose + self._refresh_target_joints_from_groups() + + def _prune_inactive_group_state(self) -> None: + selected = set(self.state.selected_group_ids) + for values in ( + self.state.pose_targets, + self.state.group_joint_targets, + self.state.group_poses, + ): + for group_id in tuple(values): + if group_id not in selected: + values.pop(group_id) + self._refresh_target_joints_from_groups() + + def _refresh_target_joints_from_groups(self) -> None: + names: list[str] = [] + positions: list[float] = [] + for group_id in self.state.selected_group_ids: + target = self.state.group_joint_targets.get(group_id) + if target is not None: + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + self.state.target_joints = ( + JointState({"name": names, "position": positions}) if names else None + ) + + def _active_pose_targets(self) -> dict[PlanningGroupID, Pose]: + return { + group_id: self.state.pose_targets[group_id] + for group_id in self.state.selected_group_ids + if group_id in self.state.pose_targets + } + + def _preset_values_by_local_name(self, preset: str, robot_name: str) -> dict[str, float]: + if preset == "Current": + state = self.get_current_joint_state(robot_name) + elif preset == "Init": + state = self.get_init_joints(robot_name) + else: + config = self.get_robot_config(robot_name) + if config is None: + return {} + return { + str(name): float(value) + for name, value in zip(config.joint_names, config.home_joints or [], strict=False) + } + return self._local_values_for_robot(robot_name, state) + def _remove_panel_handles(self) -> None: for key, handle in list(self._handles.items()): remove = getattr(handle, "remove", None) @@ -284,51 +677,20 @@ def _remove_panel_handles(self) -> None: remove() self._handles.pop(key, None) - def _select_robot(self, robot_name: str) -> None: - if self._closed: - return - if (robot_name or None) == self.state.selected_robot: - self.refresh() - return - self.state.selected_robot = robot_name or None - self.state.target_status = TargetStatus.EMPTY - self.state.feasibility.status = FeasibilityStatus.UNKNOWN - self.state.plan_state = PanelPlanState() - self._build_joint_sliders() - self._sync_preset_dropdown() - self.refresh() - - def _sync_robot_dropdown(self, robots: list[str]) -> None: - handle = self._handles.get("robot") - if handle is None: - return - options = robots or [""] - for attr in ("options", "values"): - if hasattr(handle, attr): - try: - self._set_optional_handle_attr(handle, attr, options) - except Exception: - logger.warning("Could not set robot dropdown %s", attr, exc_info=True) - if hasattr(handle, "value") and self.state.selected_robot in robots: - try: - self._set_optional_handle_attr(handle, "value", self.state.selected_robot) - except Exception: - logger.warning("Could not set robot dropdown value", exc_info=True) - def _sync_preset_dropdown(self) -> None: handle = self._handles.get("preset") - if handle is None or self.state.selected_robot is None: + if handle is None or not self.state.selected_group_ids: return - info: RobotInfo | None = self.adapter.get_robot_info(self.state.selected_robot) - config = self.adapter.get_robot_config(self.state.selected_robot) options = ["Select preset..."] - if (info is not None and info["init_joints"] is not None) or self.adapter.get_init_joints( - self.state.selected_robot - ) is not None: + selected_robots = self._selected_robot_names() + if any(self.get_init_joints(robot_name) is not None for robot_name in selected_robots): options.append("Init") options.append("Current") - home_joints = config.home_joints if config is not None else None - if (info is not None and info["home_joints"] is not None) or home_joints is not None: + if any( + (config := self.get_robot_config(robot_name)) is not None + and config.home_joints is not None + for robot_name in selected_robots + ): options.append("Home") for attr in ("options", "values"): if hasattr(handle, attr): @@ -340,253 +702,409 @@ def _sync_preset_dropdown(self) -> None: def _apply_preset(self, preset: str) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return - config = self.adapter.get_robot_config(robot_name) - if config is None: - return - if preset == "Current": - current = self.adapter.get_current_joint_state(robot_name) - values = list(current.position) if current is not None else [] - elif preset == "Init": - init = self.adapter.get_init_joints(robot_name) - values = list(init.position) if init is not None else [] - elif preset == "Home": - values = list(config.home_joints or []) - else: + if preset not in {"Current", "Init", "Home"}: return - self._set_slider_values(config.joint_names, values) - self.state.joint_target = [float(value) for value in values] + targets: dict[PlanningGroupID, JointState] = {} + slider_values: list[tuple[PlanningGroupID, tuple[str, ...], list[float]]] = [] + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + self._set_recoverable_error(f"Unknown planning group: {group_id}") + return + if preset == "Current" and self.is_state_stale(group.robot_name): + self._set_recoverable_error( + f"Cannot apply Current preset without fresh telemetry for: {group.robot_name}" + ) + return + values = self._preset_values_by_local_name(preset, str(group.robot_name)) + missing = [str(name) for name in group.local_joint_names if str(name) not in values] + if missing: + self._set_recoverable_error( + f"Cannot apply {preset} preset: missing joints for {group_id}: {', '.join(missing)}" + ) + return + positions = [float(values[str(name)]) for name in group.local_joint_names] + targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) + slider_values.append((group_id, group.local_joint_names, positions)) + self.state.group_joint_targets.update(targets) + if any( + (group_id, str(local_name)) not in self._joint_sliders + for group_id, local_names, _positions in slider_values + for local_name in local_names + ): + self._build_joint_sliders() + for group_id, local_names, positions in slider_values: + self._set_group_slider_values(group_id, local_names, positions) + self._refresh_target_joints_from_groups() self._submit_joint_target_evaluation() self.refresh() - def _set_slider_values(self, joint_names: list[str], values: list[float]) -> None: + def _set_group_slider_values( + self, group_id: PlanningGroupID, local_names: tuple[str, ...], values: list[float] + ) -> None: self._suppress_target_callbacks = True try: - for joint_name, value in zip(joint_names, values, strict=False): - handle = self._joint_sliders.get(joint_name) + for local_name, value in zip(local_names, values, strict=True): + handle = self._joint_sliders.get((group_id, str(local_name))) if handle is not None: handle.value = float(value) finally: self._suppress_target_callbacks = False - def _target_from_sliders(self, robot_name: str) -> JointState | None: - config = self.adapter.get_robot_config(robot_name) - if config is None: - self._set_error("No robot config") - return None - values = [ - float(self._joint_sliders[name].value) - for name in config.joint_names - if name in self._joint_sliders - ] - return self.adapter.joints_from_values(config.joint_names, values) - - def _on_joint_slider_update(self, _joint_name: str) -> None: + def _target_set_from_sliders(self) -> dict[PlanningGroupID, JointState] | None: + targets: dict[PlanningGroupID, JointState] = {} + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + self._set_error(f"Unknown planning group: {group_id}") + return None + positions: list[float] = [] + for local_name in group.local_joint_names: + handle = self._joint_sliders.get((group_id, str(local_name))) + if handle is None: + self._set_error(f"Missing target slider for {group_id}/{local_name}") + return None + positions.append(float(handle.value)) + targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) + return targets + + def _on_joint_slider_update(self, _group_id: PlanningGroupID, _local_name: str) -> None: if self._closed: return if self._suppress_target_callbacks: return self._submit_joint_target_evaluation() - def _on_transform_update(self, target: TransformControlsHandle) -> None: + def _on_transform_update( + self, group_id: PlanningGroupID, target: TransformControlsHandle + ) -> None: if self._closed: return - if self._suppress_target_callbacks or self.state.selected_robot is None: + if self._suppress_target_callbacks or group_id not in self.state.selected_group_ids: return pose = self._pose_from_transform_target(target) if pose is None: return self.state.cartesian_target = pose + self.state.pose_targets[group_id] = pose sequence_id = self.state.next_sequence_id() self._worker.submit( TargetEvaluationRequest( sequence_id=sequence_id, source="cartesian", - robot_name=self.state.selected_robot, - pose=pose, + selection_epoch=self.state.selection_epoch, + group_ids=self.state.selected_group_ids, + auxiliary_group_ids=tuple( + selected_group_id + for selected_group_id in self.state.selected_group_ids + if selected_group_id not in self._active_pose_targets() + ), + joints=( + None + if self.state.target_joints is None + else JointState(self.state.target_joints) + ), + pose_targets=dict(self._active_pose_targets()), ) ) self.refresh() def _submit_joint_target_evaluation(self) -> None: - robot_name = self.state.selected_robot - if robot_name is None: - return - target = self._target_from_sliders(robot_name) - if target is None: + targets = self._target_set_from_sliders() + if targets is None: return - self.state.joint_target = list(target.position) - self._move_joint_target_visuals(robot_name, target) + self.state.group_joint_targets = targets + self._refresh_target_joints_from_groups() + self._move_joint_target_visuals(targets) sequence_id = self.state.next_sequence_id() self._worker.submit( TargetEvaluationRequest( sequence_id=sequence_id, source="joints", - robot_name=robot_name, - joints=target, + selection_epoch=self.state.selection_epoch, + group_ids=self.state.selected_group_ids, + joint_targets=dict(targets), ) ) self.refresh() - def _move_joint_target_visuals(self, robot_name: str, target: JointState) -> None: + def _move_joint_target_visuals(self, targets: Mapping[PlanningGroupID, JointState]) -> None: """Optimistically move target visuals before collision/feasibility returns.""" - config = self.adapter.get_robot_config(robot_name) - robot_id = self.adapter.robot_id_for_name(robot_name) - if self.scene is not None and config is not None and robot_id is not None: - self.scene.set_target_joints(str(robot_id), config.joint_names, list(target.position)) - pose = self.adapter.get_ee_pose(robot_name, target) - if pose is not None: - self._suppress_target_callbacks = True - try: - self.scene.set_target_pose(str(robot_id), pose) - finally: - self._suppress_target_callbacks = False + if self.scene is None: + return + for robot_name, state in self._target_ghost_states(targets).items(): + config = self.get_robot_config(robot_name) + robot_id = self.robot_id_for_name(robot_name) + if config is not None and robot_id is not None: + self.scene.set_target_joints(str(robot_id), config.joint_names, state.position) + + def _target_ghost_states( + self, targets: Mapping[PlanningGroupID, JointState] + ) -> dict[str, JointState]: + groups = self._groups_by_id() + merged: dict[str, dict[str, float]] = {} + configs: dict[str, tuple[str, ...]] = {} + for group_id in self.state.selected_group_ids: + group = groups.get(group_id) + target = targets.get(group_id) + if group is None or target is None: + continue + robot_name = str(group.robot_name) + config = self.get_robot_config(robot_name) + current = self.get_current_joint_state(robot_name) + if config is None or current is None: + continue + values = self._local_values_for_robot(robot_name, current) + target_raw = self._state_values_by_local_name(target) + for local_name, global_name in zip( + group.local_joint_names, group.joint_names, strict=True + ): + if str(global_name) in target_raw: + values[str(local_name)] = target_raw[str(global_name)] + elif str(local_name) in target_raw: + values[str(local_name)] = target_raw[str(local_name)] + if all(name in values for name in config.joint_names): + merged[robot_name] = values + configs[robot_name] = tuple(config.joint_names) + return { + robot_name: JointState( + {"name": list(joint_names), "position": [values[name] for name in joint_names]} + ) + for robot_name, values in merged.items() + for joint_names in (configs[robot_name],) + } + + def _sync_target_ghost_visibility(self) -> None: + if self.scene is None: + return + active_robot_ids = { + str(robot_id) + for group_id in self.state.selected_group_ids + if (group := self._groups_by_id().get(group_id)) is not None + and group.has_pose_target + and (robot_id := self.robot_id_for_name(group.robot_name)) is not None + } + for _robot_name, robot_id, _config in self.robot_items(): + self.scene.set_target_active(str(robot_id), str(robot_id) in active_robot_ids) def _handle_target_evaluation_request( self, request: TargetEvaluationRequest - ) -> TargetEvaluation: + ) -> TargetEvaluationResult: if request.source == "cartesian": - if request.pose is None: - return {"success": False, "status": "INVALID", "message": "No pose target"} - return self.adapter.evaluate_pose_target(request.pose, request.robot_name) - if request.joints is None: - return {"success": False, "status": "INVALID", "message": "No joint target"} - return self.adapter.evaluate_joint_target(request.joints, request.robot_name) + if not request.pose_targets: + return TargetEvaluationResult(False, "INVALID", "No pose target") + return self.evaluate_pose_target_set( + request.pose_targets, request.auxiliary_group_ids, request.joints + ) + if not request.joint_targets: + return TargetEvaluationResult(False, "INVALID", "No joint target") + return self.evaluate_joint_target_set(request.group_ids, request.joint_targets) def _apply_target_evaluation_result( - self, request: TargetEvaluationRequest, result: TargetEvaluation + self, request: TargetEvaluationRequest, result: TargetEvaluationResult ) -> None: if self._closed: return - if request.sequence_id != self.state.latest_sequence_id: + if ( + request.sequence_id != self.state.latest_sequence_id + or request.selection_epoch != self.state.selection_epoch + or request.group_ids != self.state.selected_group_ids + ): return - collision_free = bool(result.get("collision_free", False)) - success = bool(result.get("success", False)) + collision_free = result.collision_free + success = result.success self.state.feasibility.status = self._feasibility_status(result, success, collision_free) - self.state.feasibility.message = str(result.get("message", "")) + self.state.feasibility.message = result.message self.state.target_status = ( TargetStatus.FEASIBLE if success and collision_free else TargetStatus.INFEASIBLE ) self.state.error = "" if success and collision_free else self.state.feasibility.message + if result.target_joints is not None: + self.state.target_joints = JointState(result.target_joints) + self._split_target_joints_by_group(result.target_joints) + self.state.group_poses = { + str(group_id): pose + for group_id, pose in result.group_poses.items() + if isinstance(pose, Pose) + } if request.source == "joints": - joint_state = result.get("joint_state") - if isinstance(joint_state, JointState): - self.state.joint_target = list(joint_state.position) - if request.source == "cartesian": - joint_state = result.get("joint_state") - if isinstance(joint_state, JointState): - self.state.joint_target = list(joint_state.position) - pose = result.get("ee_pose") - if isinstance(pose, Pose): - self.state.cartesian_target = pose + self._sync_pose_targets_from_group_poses() + else: self._sync_controls_from_targets() self._update_target_visual_state() self.refresh() def _sync_controls_from_targets(self) -> None: - robot_name = self.state.selected_robot - if robot_name is None: - return - config = self.adapter.get_robot_config(robot_name) - if config is not None and self.state.joint_target is not None: - self._set_slider_values(list(config.joint_names), list(self.state.joint_target)) - robot_id = self.adapter.robot_id_for_name(robot_name) - if self.scene is not None and robot_id is not None: - self.scene.set_target_joints( - str(robot_id), config.joint_names, self.state.joint_target + for group_id, target in self.state.group_joint_targets.items(): + group = self._groups_by_id().get(group_id) + if group is not None: + self._set_group_slider_values( + group_id, group.local_joint_names, list(target.position) ) - # Do not write the Cartesian target back into the active transform - # control here. The gizmo is the source of truth for Cartesian edits; - # programmatic pose writes from delayed IK results can fight fast user - # dragging and make the gizmo jump back. + self._move_joint_target_visuals(self.state.group_joint_targets) + + def _split_target_joints_by_group(self, target_joints: JointState) -> None: + if len(target_joints.name) != len(target_joints.position): + return + positions = { + str(name): float(value) + for name, value in zip(target_joints.name, target_joints.position, strict=True) + } + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None or any(str(name) not in positions for name in group.joint_names): + continue + self.state.group_joint_targets[group_id] = JointState( + { + "name": list(group.joint_names), + "position": [positions[str(name)] for name in group.joint_names], + } + ) + + def _sync_pose_targets_from_group_poses(self) -> None: + groups = self._groups_by_id() + active_group_ids: list[PlanningGroupID] = [] + for group_id, pose in self.state.group_poses.items(): + group = groups.get(group_id) + if group is None or not group.has_pose_target: + continue + if group_id not in self.state.selected_group_ids: + continue + self.state.pose_targets[group_id] = pose + active_group_ids.append(group_id) + if self.scene is None: + return + self._suppress_target_callbacks = True + try: + for group_id in active_group_ids: + self.scene.set_target_pose(str(group_id), self.state.pose_targets[group_id]) + finally: + self._suppress_target_callbacks = False def _update_status_text(self) -> None: current = self.state.current_joints + status_label = self.state.error or self.state.module_state status = [ - "### Manipulation Panel", - f"Robot: `{self.state.selected_robot or 'none'}`", - f"Module: `{self.state.module_state}`", - f"Backend: `{self.state.backend_status.value}`", - f"Target: `{self.state.target_status.value}`", - f"Feasibility: `{self.state.feasibility.status.value}`", - f"Plan: `{self.state.plan_state.status.value}`", - f"Action: `{self.state.action_status.value}`", + "### Status", + f"**State:** {status_label}", + f"Target: `{self.state.target_status.value}` · Plan: `{self.state.plan_state.status.value}`", ] - if self.state.selected_robot is not None: - status.append( - f"State stale: `{self.adapter.is_state_stale(self.state.selected_robot)}`" - ) + stale_robots = self._stale_robot_names(self.state.selected_group_ids) + if self.state.selected_group_ids: + stale_detail = "False" if not stale_robots else f"True ({', '.join(stale_robots)})" + status.append(f"State stale: `{stale_detail}`") if current is not None: status.append(f"Current joints: `{[round(v, 3) for v in current]}`") if self.state.last_result: status.append(f"Last result: `{self.state.last_result}`") - if self.state.error: - status.append(f"Error: `{self.state.error}`") self._set_handle_value("status", "\n\n".join(status)) + self._set_handle_value( + "target_summary", + f"Feasibility: `{self.state.feasibility.status.value}`", + ) def _update_control_state(self) -> None: self._set_disabled("plan", not self.state.can_plan()) self._set_disabled("preview", not self.state.can_preview()) self._set_disabled( "execute", - not self.state.can_execute(self.config.current_match_tolerance), + not self._can_execute(), ) - self._set_disabled("cancel", not self.state.can_cancel()) + can_cancel = self.state.can_cancel() + self._set_disabled("cancel", not can_cancel) + self._set_visible("cancel", can_cancel) self._update_target_visual_state() def _update_target_visual_state(self) -> None: - if self.scene is None or self.state.selected_robot is None: - return - robot_id = self.adapter.robot_id_for_name(self.state.selected_robot) - if robot_id is None: + if self.scene is None: return - self.scene.set_target_visual_state( - str(robot_id), self.state.feasibility.status == FeasibilityStatus.FEASIBLE + feasible = self.state.feasibility.status == FeasibilityStatus.FEASIBLE + groups = self._groups_by_id() + selected_groups = tuple( + (group_id, groups[group_id]) + for group_id in self.state.selected_group_ids + if group_id in groups + ) + for group_id, group in selected_groups: + if group.has_pose_target: + self.scene.set_target_control_visual_state(str(group_id), feasible) + robot_ids = tuple( + dict.fromkeys( + str(robot_id) + for _group_id, group in selected_groups + if (robot_id := self.robot_id_for_name(str(group.robot_name))) is not None + ) ) + for robot_id in robot_ids: + self.scene.set_target_robot_visual_state(robot_id, feasible) + + def _can_execute(self) -> bool: + return self.state.can_execute() def _submit_plan(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return if not self.state.can_plan(): self._set_recoverable_error( "Cannot plan until target is feasible and manipulation is idle" ) return + group_ids = self.state.selected_group_ids + selection_epoch = self.state.selection_epoch + target_sequence_id = self.state.latest_sequence_id + targets = self._target_set_from_sliders() + if targets is None: + return operation_id = self._next_operation_id() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): return self.state.action_status = ActionStatus.RUNNING self.state.plan_state.status = PlanStatus.PLANNING - if self.state.manipulation_state == "FAULT" and not self.adapter.reset(): - self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation("reset=False", clear_error=False, operation_id=operation_id) + stale_robots = self._stale_robot_names(group_ids) + if stale_robots: + if not self._operation_is_current( + operation_id, selection_epoch, target_sequence_id + ): + self._finish_operation( + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch + ) + return + self.state.plan_state.status = PlanStatus.STALE + self.state.error = "Cannot plan without fresh telemetry for: " + ", ".join( + stale_robots + ) + self._finish_operation( + "plan=False", + clear_error=False, + operation_id=operation_id, + selection_epoch=selection_epoch, + ) return - target = self._target_from_sliders(robot_name) - if target is None: - self.state.plan_state.status = PlanStatus.FAILED + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): self._finish_operation( - "plan_to_joints=False", clear_error=False, operation_id=operation_id + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch ) return - ok = self.adapter.plan_to_joints(target, robot_name) - if not self._operation_is_current(operation_id): + ok = self.plan_to_selected_joints(group_ids, targets) + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return if ok: self.state.plan_state.status = PlanStatus.FRESH - self.state.plan_state.robot = robot_name - self.state.plan_state.target_joints = list(target.position) - self.state.plan_state.target_pose = self.state.cartesian_target - self.state.plan_state.start_joints_snapshot = list(self.state.current_joints or []) + self.state.plan_state.group_ids = group_ids + self.state.plan_state.target_sequence_id = target_sequence_id else: self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation(f"plan_to_joints={ok}", operation_id=operation_id) + self.state.plan_state.plan = None + self._finish_operation( + f"plan_to_joints={ok}", + operation_id=operation_id, + selection_epoch=selection_epoch, + ) self._operation_worker.submit( operation, on_error=lambda message: self._set_operation_error(message, operation_id) @@ -595,20 +1113,21 @@ def operation() -> None: def _submit_preview(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return if not self.state.can_preview(): self._set_recoverable_error("No fresh plan to preview") return + selection_epoch = self.state.selection_epoch operation_id = self._next_operation_id() + self.state.action_status = ActionStatus.PREVIEWING + self.refresh() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch): return - self.state.action_status = ActionStatus.PREVIEWING - ok = self.adapter.preview_path(robot_name) - self._finish_operation(f"preview={ok}", operation_id=operation_id) + ok = self.preview_path() + self._finish_operation( + f"preview={ok}", operation_id=operation_id, selection_epoch=selection_epoch + ) self._operation_worker.submit( operation, @@ -616,30 +1135,56 @@ def operation() -> None: on_error=lambda message: self._set_operation_error(message, operation_id), ) + def _clear_invalidated_preview(self) -> None: + if self.state.action_status == ActionStatus.PREVIEWING: + self._operation_sequence_id += 1 + self.state.action_status = ActionStatus.IDLE + self.state.last_result = "preview=False" + def _submit_execute(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return - if not self.state.can_execute(self.config.current_match_tolerance): - self._set_recoverable_error( - "Cannot execute: require feasible fresh plan and matching current joints" - ) + if not self._can_execute(): + self._set_recoverable_error("Cannot execute: require feasible fresh plan") return + selection_epoch = self.state.selection_epoch + group_ids = self.state.selected_group_ids + target_sequence_id = self.state.latest_sequence_id operation_id = self._next_operation_id() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + return + if ( + self.state.plan_state.group_ids != group_ids + or self.state.plan_state.target_sequence_id != target_sequence_id + ): + self.state.plan_state.status = PlanStatus.STALE + self._finish_operation( + "execute=False", + clear_error=False, + operation_id=operation_id, + selection_epoch=selection_epoch, + ) + return + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "execute=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return self.state.action_status = ActionStatus.EXECUTING self.state.plan_state.status = PlanStatus.EXECUTING - ok = self.adapter.execute(robot_name) - if not self._operation_is_current(operation_id): + ok = self.execute() + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "execute=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return if not ok: self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation(f"execute={ok}", operation_id=operation_id) + self._finish_operation( + f"execute={ok}", operation_id=operation_id, selection_epoch=selection_epoch + ) self._operation_worker.submit( operation, on_error=lambda message: self._set_operation_error(message, operation_id) @@ -656,7 +1201,7 @@ def _submit_cancel(self) -> None: self._mark_cancelled_plan_state(cancelled_action) self._restart_operation_worker() try: - ok = self.adapter.cancel() + ok = self.cancel() except Exception as e: self._set_operation_error(str(e), operation_id) return @@ -685,7 +1230,7 @@ def operation() -> None: if not self._operation_is_current(operation_id): return self.state.action_status = ActionStatus.CLEARING_PLAN - ok = self.adapter.clear_planned_path() + ok = self.clear_planned_path() if not self._operation_is_current(operation_id): return self.state.plan_state = PanelPlanState() @@ -699,14 +1244,30 @@ def _next_operation_id(self) -> int: self._operation_sequence_id += 1 return self._operation_sequence_id - def _operation_is_current(self, operation_id: int) -> bool: - return not self._closed and operation_id == self._operation_sequence_id + def _operation_is_current( + self, + operation_id: int, + selection_epoch: int | None = None, + target_sequence_id: int | None = None, + ) -> bool: + return ( + not self._closed + and operation_id == self._operation_sequence_id + and (selection_epoch is None or selection_epoch == self.state.selection_epoch) + and (target_sequence_id is None or target_sequence_id == self.state.latest_sequence_id) + ) def _finish_operation( - self, result: str, *, clear_error: bool = True, operation_id: int | None = None + self, + result: str, + *, + clear_error: bool = True, + operation_id: int | None = None, + selection_epoch: int | None = None, ) -> None: if self._closed or ( - operation_id is not None and not self._operation_is_current(operation_id) + operation_id is not None + and not self._operation_is_current(operation_id, selection_epoch) ): return self.state.action_status = ActionStatus.IDLE @@ -743,6 +1304,11 @@ def _set_disabled(self, key: str, disabled: bool) -> None: if isinstance(handle, GuiButtonHandle): self._set_optional_handle_attr(handle, "disabled", disabled) + def _set_visible(self, key: str, visible: bool) -> None: + handle = self._handles.get(key) + if handle is not None: + self._set_optional_handle_attr(handle, "visible", visible) + @staticmethod def _set_optional_handle_attr(handle: object, attr: str, value: object) -> None: setattr(handle, attr, value) @@ -753,9 +1319,9 @@ def _pose_from_transform_target(self, target: TransformControlsHandle) -> Pose | return Pose({"position": [px, py, pz], "orientation": [qx, qy, qz, qw]}) def _feasibility_status( - self, result: TargetEvaluation, success: bool, collision_free: bool + self, result: TargetEvaluationResult, success: bool, collision_free: bool ) -> FeasibilityStatus: - status = str(result.get("status", "")).upper() + status = result.status.upper() if success and collision_free: return FeasibilityStatus.FEASIBLE if status in {"COLLISION", "COLLISION_AT_START", "COLLISION_AT_GOAL"}: diff --git a/dimos/manipulation/visualization/viser/scene.py b/dimos/manipulation/visualization/viser/scene.py index cd50ab098e..296e18319f 100644 --- a/dimos/manipulation/visualization/viser/scene.py +++ b/dimos/manipulation/visualization/viser/scene.py @@ -15,22 +15,35 @@ from __future__ import annotations from collections.abc import Callable, Sequence +import hashlib +import os from pathlib import Path +import tempfile +from threading import RLock +import time from typing import Protocol, TypeAlias, cast +import xml.etree.ElementTree as ET from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake -from dimos.manipulation.visualization.viser.animation import PreviewAnimator +from dimos.manipulation.visualization.viser.animation import ( + GroupPreviewAnimation, + PreviewFrame, + preview_tick_times, + scaled_frame_delays, +) from dimos.manipulation.visualization.viser.runtime import ( VISER_INSTALL_HINT, VISER_URDF_INSTALL_HINT, ) from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.model_parser import parse_model from dimos.utils.logging_config import setup_logger try: from viser import ( + FrameHandle, GridHandle, MeshHandle, TransformControlsEvent, @@ -55,6 +68,8 @@ logger = setup_logger() +_VISER_URDF_CACHE_DIR = Path(tempfile.gettempdir()) / "dimos_viser_urdf_cache" + GOAL_ROBOT_FEASIBLE_COLOR = (255, 122, 0) GOAL_ROBOT_INFEASIBLE_COLOR = (255, 30, 30) GOAL_ROBOT_FEASIBLE_OPACITY = 0.7 @@ -69,7 +84,7 @@ REFERENCE_GRID_CELL_COLOR = (44, 54, 58) REFERENCE_GRID_SECTION_COLOR = (90, 145, 165) -SceneHandle: TypeAlias = ViserUrdf | TransformControlsHandle | GridHandle | MeshHandle +SceneHandle: TypeAlias = ViserUrdf | TransformControlsHandle | GridHandle | MeshHandle | FrameHandle class _ColorHandle(Protocol): @@ -79,19 +94,21 @@ class _ColorHandle(Protocol): class ViserManipulationScene: """Viser scene graph helpers for current robot, ghost robot, and path rendering.""" - def __init__( - self, server: ViserServer, viser_urdf: type[ViserUrdf], *, preview_fps: float - ) -> None: + def __init__(self, server: ViserServer, viser_urdf: type[ViserUrdf]) -> None: self.server = server self.viser_urdf = viser_urdf - self.preview_fps = preview_fps self._configs_by_id: dict[str, RobotModelConfig] = {} self._urdfs: dict[str, ViserUrdf] = {} self._handles: dict[str, TransformControlsHandle] = {} + self._root_frames: dict[str, FrameHandle] = {} self._grid_handle: GridHandle | None = None self._grid_visible = True self._preview_visible: dict[str, bool] = {} + self._target_active: dict[str, bool] = {} self._target_tracks_current: dict[str, bool] = {} + self._scene_lock = RLock() + self._animation_generation = 0 + self._animation_generations: dict[str, int] = {} self._ensure_reference_grid() def has_reference_grid(self) -> bool: @@ -106,9 +123,18 @@ def set_reference_grid_visible(self, visible: bool) -> None: def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: self._configs_by_id[robot_id] = config self._preview_visible.setdefault(robot_id, False) + self._animation_generations.setdefault(robot_id, 0) + self._target_active.setdefault(robot_id, False) self._target_tracks_current.setdefault(robot_id, True) self._ensure_robot_urdfs(robot_id, config) + def set_target_active(self, robot_id: str, active: bool) -> None: + """Show the target ghost only while a pose-target group is selected.""" + self._target_active[robot_id] = active + if not active: + self._target_tracks_current[robot_id] = True + self._set_target_visibility(robot_id, active) + def _ensure_reference_grid(self) -> None: try: scene = self.server.scene @@ -155,41 +181,104 @@ def dispatch(event: TransformControlsEvent) -> None: self._handles[handle_key] = handle return handle + def remove_target_controls(self, control_id: str) -> None: + self._remove_handle(f"{control_id}:ee_control") + def update_current_robot(self, robot_id: str, joint_state: JointState | None) -> None: - config = self._configs_by_id.get(robot_id) - if config is None or joint_state is None: - return - self._ensure_robot_urdfs(robot_id, config) - current = self._urdfs.get(f"{robot_id}:current") - self.set_urdf_joints(current, config.joint_names, joint_state.position) - if self._target_tracks_current.get(robot_id, True): - self._set_target_joints(robot_id, config.joint_names, joint_state.position) - self._set_target_visibility(robot_id, True) + with self._scene_lock: + config = self._configs_by_id.get(robot_id) + if config is None or joint_state is None: + return + self._ensure_robot_urdfs(robot_id, config) + current = self._urdfs.get(f"{robot_id}:current") + self.set_urdf_joints(current, config.joint_names, joint_state.position) + if self._target_tracks_current.get(robot_id, True): + self._set_target_joints(robot_id, config.joint_names, joint_state.position) + self._set_target_visibility(robot_id, self._target_active.get(robot_id, False)) + + def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> None: + """Prevent an old blocking animation from touching replacement handles.""" + with self._scene_lock: + self._animation_generation += 1 + affected = set(robot_ids) if robot_ids is not None else set(self._preview_visible) + for robot_id in affected: + self._animation_generations[robot_id] = ( + self._animation_generations.get(robot_id, 0) + 1 + ) + if robot_id not in self._preview_visible: + continue + self._preview_visible[robot_id] = False + self._set_preview_visibility(robot_id, False) - def show_preview(self, robot_id: str) -> None: - """Show the transient preview-animation ghost. + def animate_preview(self, preview: GroupPreviewAnimation, duration: float) -> bool: + """Play every robot from one normalized tick clock. - Target editing uses the separate target ghost and must not call this path. + Inputs are fully validated before ghosts become visible; a generation + replacement, clear, or close stops mutation before the next tick. """ - self._preview_visible[robot_id] = True - self._set_preview_visibility(robot_id, True) - - def hide_preview(self, robot_id: str) -> None: - """Hide the transient preview-animation ghost.""" - self._preview_visible[robot_id] = False - self._set_preview_visibility(robot_id, False) - - def animate_path(self, robot_id: str, path: Sequence[JointState], duration: float) -> bool: - config = self._configs_by_id.get(robot_id) - if config is None: + frames = {track.robot_id: track.frames for track in preview.tracks} + names = {track.robot_id: track.joint_names for track in preview.tracks} + if ( + not frames + or len(frames) != len(preview.tracks) + or any( + not values or robot_id not in self._configs_by_id + for robot_id, values in frames.items() + ) + ): return False - self.show_preview(robot_id) + tick_times = preview_tick_times(preview) + if not tick_times: + return False + with self._scene_lock: + self._animation_generation += 1 + generations: dict[str, int] = {} + for robot_id in frames: + self._animation_generations[robot_id] = ( + self._animation_generations.get(robot_id, 0) + 1 + ) + generations[robot_id] = self._animation_generations[robot_id] + self._preview_visible[robot_id] = True + self._set_preview_visibility(robot_id, True) try: - return PreviewAnimator( - lambda joints: self._set_preview_ghost_joints(robot_id, config.joint_names, joints) - ).animate(path, duration, self.preview_fps) + delays = scaled_frame_delays( + tuple( + PreviewFrame(time_from_start=tick_time, positions=()) + for tick_time in tick_times + ), + duration, + ) + frame_indices = {robot_id: 0 for robot_id in frames} + for tick, tick_time in enumerate(tick_times): + with self._scene_lock: + active_robot_ids = [ + robot_id + for robot_id in frames + if self._animation_generations.get(robot_id) == generations[robot_id] + ] + if not active_robot_ids: + return False + for robot_id in active_robot_ids: + robot_frames = frames[robot_id] + while ( + frame_indices[robot_id] + 1 < len(robot_frames) + and robot_frames[frame_indices[robot_id] + 1].time_from_start + <= tick_time + ): + frame_indices[robot_id] += 1 + source = frame_indices[robot_id] + self._set_preview_ghost_joints( + robot_id, names[robot_id], robot_frames[source].positions + ) + if tick < len(delays): + time.sleep(delays[tick]) + return True finally: - self.hide_preview(robot_id) + with self._scene_lock: + for robot_id in frames: + if self._animation_generations.get(robot_id) == generations[robot_id]: + self._preview_visible[robot_id] = False + self._set_preview_visibility(robot_id, False) def set_target_joints( self, robot_id: str, joint_names: Sequence[str], joints: Sequence[float] @@ -235,16 +324,26 @@ def set_target_pose(self, robot_id: str, pose: Pose | None) -> None: ) def set_target_visual_state(self, robot_id: str, feasible: bool) -> None: + """Set the legacy matching robot/control target visual state.""" + self.set_target_control_visual_state(robot_id, feasible) + self.set_target_robot_visual_state(robot_id, feasible) + + def set_target_control_visual_state(self, control_id: str, feasible: bool) -> None: + """Set feasibility color for one planning-group keyed target control.""" color = TARGET_CONTROL_FEASIBLE_COLOR if feasible else TARGET_CONTROL_INFEASIBLE_COLOR - mesh_color = GOAL_ROBOT_FEASIBLE_COLOR if feasible else GOAL_ROBOT_INFEASIBLE_COLOR - mesh_opacity = GOAL_ROBOT_FEASIBLE_OPACITY if feasible else GOAL_ROBOT_INFEASIBLE_OPACITY - handle = self._handles.get(f"{robot_id}:ee_control") + handle = self._handles.get(f"{control_id}:ee_control") if handle is not None: cast("_ColorHandle", handle).color = color + + def set_target_robot_visual_state(self, robot_id: str, feasible: bool) -> None: + """Set feasibility material for one robot-ID keyed target ghost.""" + mesh_color = GOAL_ROBOT_FEASIBLE_COLOR if feasible else GOAL_ROBOT_INFEASIBLE_COLOR + mesh_opacity = GOAL_ROBOT_FEASIBLE_OPACITY if feasible else GOAL_ROBOT_INFEASIBLE_OPACITY target = self._urdfs.get(f"{robot_id}:target") self._set_urdf_mesh_material(target, mesh_color, mesh_opacity) def close(self) -> None: + self.cancel_preview_animation() for key in list(self._handles): self._remove_handle(key) if self._grid_handle is not None: @@ -252,9 +351,13 @@ def close(self) -> None: self._grid_handle = None for urdf in self._urdfs.values(): self._remove_scene_handle(urdf) + for frame in self._root_frames.values(): + self._remove_scene_handle(frame) self._urdfs.clear() + self._root_frames.clear() self._configs_by_id.clear() self._preview_visible.clear() + self._target_active.clear() self._target_tracks_current.clear() def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: @@ -264,11 +367,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: key = f"{robot_id}:{kind}" if key in self._urdfs: continue - root_node_name = { - "current": f"/robots/{robot_id}/current", - "target": f"/targets/{robot_id}/target", - "preview": f"/previews/{robot_id}/ghost", - }[kind] + root_node_name = self._urdf_root_node_name(robot_id, kind, config) mesh_color_override = { "current": None, "target": GOAL_ROBOT_MESH_COLOR, @@ -284,7 +383,9 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: self._set_urdf_mesh_material( self._urdfs[key], GOAL_ROBOT_FEASIBLE_COLOR, GOAL_ROBOT_FEASIBLE_OPACITY ) - self._set_handle_visibility(self._urdfs[key], True) + self._set_handle_visibility( + self._urdfs[key], self._target_active.get(robot_id, False) + ) elif kind == "preview": self._set_urdf_mesh_material( self._urdfs[key], PREVIEW_ROBOT_COLOR, PREVIEW_ROBOT_OPACITY @@ -295,7 +396,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: def prepared_urdf_path(self, config: RobotModelConfig) -> Path: package_paths = {package: Path(path) for package, path in config.package_paths.items()} - return Path( + prepared_path = Path( prepare_urdf_for_drake( Path(str(config.model_path)), package_paths=package_paths, @@ -303,6 +404,137 @@ def prepared_urdf_path(self, config: RobotModelConfig) -> Path: convert_meshes=bool(config.auto_convert_meshes), ) ) + prepared_path = self._strip_visualization_world_root_attachment(config, prepared_path) + self._assert_base_link_is_urdf_root(config, prepared_path) + return prepared_path + + @staticmethod + def _strip_visualization_world_root_attachment( + config: RobotModelConfig, prepared_path: Path + ) -> Path: + """Detach a model-owned world root only for Viser base-pose rendering.""" + urdf_content = prepared_path.read_text() + try: + root = ET.fromstring(urdf_content) + except ET.ParseError: + return prepared_path + + attachments = [ + joint + for joint in root.findall("joint") + if joint.attrib.get("type") == "fixed" + and (parent := joint.find("parent")) is not None + and parent.attrib.get("link") == "world" + and (child := joint.find("child")) is not None + and child.attrib.get("link") == config.base_link + ] + if len(attachments) != 1: + return prepared_path + + root.remove(attachments[0]) + referenced_links = { + link + for joint in root.findall("joint") + for element in (joint.find("parent"), joint.find("child")) + if (link := element.get("link") if element is not None else None) is not None + } + if "world" in referenced_links: + return prepared_path + world_links = [link for link in root.findall("link") if link.attrib.get("name") == "world"] + if len(world_links) != 1: + return prepared_path + root.remove(world_links[0]) + + stripped_content = ET.tostring(root, encoding="unicode") + digest = hashlib.sha256( + f"viser-world-root-v1\0{config.base_link}\0{urdf_content}".encode() + ).hexdigest() + cache_path = _VISER_URDF_CACHE_DIR / f"{digest}.urdf" + if cache_path.exists(): + return cache_path + _VISER_URDF_CACHE_DIR.mkdir(parents=True, exist_ok=True) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=_VISER_URDF_CACHE_DIR, + prefix=f".{digest}.", + suffix=".tmp", + delete=False, + ) as temporary_file: + temporary_file.write(stripped_content) + temporary_path = Path(temporary_file.name) + os.replace(temporary_path, cache_path) + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + return cache_path + + @staticmethod + def _assert_base_link_is_urdf_root(config: RobotModelConfig, prepared_path: Path) -> None: + root_link = parse_model(prepared_path).root_link + if root_link == config.base_link: + return + raise ValueError( + f"Viser visualization requires base_link '{config.base_link}' to match " + f"the prepared URDF root '{root_link}' because base_pose is applied to the URDF root" + ) + + def _urdf_root_node_name(self, robot_id: str, kind: str, config: RobotModelConfig) -> str: + root_node_name = { + "current": f"/robots/{robot_id}/current", + "target": f"/targets/{robot_id}/target", + "preview": f"/previews/{robot_id}/ghost", + }[kind] + if not self._has_non_identity_base_pose(config): + return root_node_name + self._ensure_base_pose_frame(robot_id, kind, config) + return f"{root_node_name}/base_pose/urdf" + + def _ensure_base_pose_frame(self, robot_id: str, kind: str, config: RobotModelConfig) -> None: + key = f"{robot_id}:{kind}:base_pose" + if key in self._root_frames: + return + pose = config.base_pose + frame_name = { + "current": f"/robots/{robot_id}/current/base_pose", + "target": f"/targets/{robot_id}/target/base_pose", + "preview": f"/previews/{robot_id}/ghost/base_pose", + }[kind] + self._root_frames[key] = self.server.scene.add_frame( + frame_name, + show_axes=False, + position=( + float(pose.position.x), + float(pose.position.y), + float(pose.position.z), + ), + wxyz=( + float(pose.orientation.w), + float(pose.orientation.x), + float(pose.orientation.y), + float(pose.orientation.z), + ), + ) + + @staticmethod + def _has_non_identity_base_pose(config: RobotModelConfig) -> bool: + pose = getattr(config, "base_pose", None) + if pose is None: + return False + return any( + abs(value) > 1e-12 + for value in ( + float(pose.position.x), + float(pose.position.y), + float(pose.position.z), + float(pose.orientation.x), + float(pose.orientation.y), + float(pose.orientation.z), + float(pose.orientation.w) - 1.0, + ) + ) def set_urdf_joints( self, urdf: ViserUrdf | None, joint_names: Sequence[str], joints: Sequence[float] @@ -328,8 +560,7 @@ def viser_joint_configuration( return [] values_by_name: dict[str, float] = {} for name, value in zip(joint_names, joints, strict=False): - values_by_name[name] = float(value) - values_by_name[name.rsplit("/", 1)[-1]] = float(value) + values_by_name[str(name)] = float(value) return [values_by_name.get(name, 0.0) for name in allowed_names] def viser_actuated_joint_names(self, urdf: ViserUrdf) -> tuple[str, ...]: diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index d8500692fd..95a21e5b84 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -21,7 +21,8 @@ import threading from typing import Literal -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID +from dimos.manipulation.visualization.operator import TargetEvaluationResult from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -93,24 +94,27 @@ class FeasibilityState: class PanelPlanState: status: PlanStatus = PlanStatus.NONE robot: str | None = None - target_pose: Pose | None = None - target_joints: list[float] | None = None - start_joints_snapshot: list[float] | None = None + group_ids: tuple[PlanningGroupID, ...] = () + target_sequence_id: int = 0 + plan: GeneratedPlan | None = None @dataclass class PanelState: selected_robot: str | None = None + selected_group_ids: tuple[PlanningGroupID, ...] = () + selection_epoch: int = 0 + pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + group_joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) + target_joints: JointState | None = None + group_poses: dict[PlanningGroupID, Pose] = field(default_factory=dict) runtime: PanelRuntime = PanelRuntime.STOPPED backend_status: BackendConnectionStatus = BackendConnectionStatus.DISCONNECTED target_status: TargetStatus = TargetStatus.EMPTY action_status: ActionStatus = ActionStatus.IDLE manipulation_state: str = "DISCONNECTED" - robot_info: RobotInfo | None = None current_joints: list[float] | None = None - current_ee_pose: Pose | None = None cartesian_target: Pose | None = None - joint_target: list[float] | None = None feasibility: FeasibilityState = field(default_factory=FeasibilityState) latest_sequence_id: int = 0 plan_state: PanelPlanState = field(default_factory=PanelPlanState) @@ -124,18 +128,25 @@ def next_sequence_id(self) -> int: self.mark_plan_stale() return self.latest_sequence_id + def advance_selection_epoch(self) -> int: + """Invalidate callbacks and plans that belong to an older group selection.""" + self.selection_epoch += 1 + self.next_sequence_id() + self.plan_state = PanelPlanState() + return self.selection_epoch + def mark_plan_stale(self) -> None: - if self.plan_state.status == PlanStatus.FRESH: + if self.plan_state.status in {PlanStatus.FRESH, PlanStatus.PLANNING}: self.plan_state.status = PlanStatus.STALE def can_plan(self) -> bool: return ( self.runtime == PanelRuntime.RUNNING and self.backend_status == BackendConnectionStatus.READY - and self.selected_robot is not None + and bool(self.selected_group_ids) and self.action_status == ActionStatus.IDLE and self.target_status == TargetStatus.FEASIBLE - and self.manipulation_state in {"IDLE", "COMPLETED", "FAULT"} + and self.manipulation_state in {"IDLE", "COMPLETED"} and self.plan_state.status != PlanStatus.PLANNING ) @@ -154,11 +165,7 @@ def can_cancel(self) -> bool: ActionStatus.EXECUTING, } or (self.manipulation_state == "EXECUTING") - def can_execute( - self, - current_tolerance: float, - action_status: ActionStatus | None = None, - ) -> bool: + def can_execute(self, action_status: ActionStatus | None = None) -> bool: plan = self.plan_state effective_action_status = action_status or self.action_status if not ( @@ -168,19 +175,12 @@ def can_execute( and self.target_status == TargetStatus.FEASIBLE and self.manipulation_state in {"IDLE", "COMPLETED"} and plan.status == PlanStatus.FRESH - and plan.robot == self.selected_robot - and plan.start_joints_snapshot is not None - and self.current_joints is not None + and plan.plan is not None + and plan.group_ids == self.selected_group_ids + and plan.target_sequence_id == self.latest_sequence_id ): return False - if len(plan.start_joints_snapshot) != len(self.current_joints): - return False - return all( - abs(expected - current) <= current_tolerance - for expected, current in zip( - plan.start_joints_snapshot, self.current_joints, strict=False - ) - ) + return True @property def connected(self) -> bool: @@ -202,9 +202,14 @@ def module_state(self) -> str: class TargetEvaluationRequest: sequence_id: int source: PreviewSource - robot_name: str + robot_name: str | None = None + selection_epoch: int = 0 + group_ids: tuple[PlanningGroupID, ...] = () pose: Pose | None = None joints: JointState | None = None + auxiliary_group_ids: tuple[PlanningGroupID, ...] = () + pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) class TargetEvaluationWorker: @@ -217,8 +222,8 @@ class TargetEvaluationWorker: def __init__( self, - handler: Callable[[TargetEvaluationRequest], TargetEvaluation], - apply_result: Callable[[TargetEvaluationRequest, TargetEvaluation], None], + handler: Callable[[TargetEvaluationRequest], TargetEvaluationResult], + apply_result: Callable[[TargetEvaluationRequest, TargetEvaluationResult], None], ) -> None: self._handler = handler self._apply_result = apply_result diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py new file mode 100644 index 0000000000..e4e83f6448 --- /dev/null +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -0,0 +1,474 @@ +# 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. + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import pytest + +pytest.importorskip("viser", reason="Viser optional dependency is not installed") + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningSceneInfo +from dimos.manipulation.visualization.operator import OperatorStatus, TargetEvaluationResult +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.manipulation.visualization.viser.gui import ViserPanelGui +from dimos.manipulation.visualization.viser.state import ( + ActionStatus, + BackendConnectionStatus, + FeasibilityStatus, + OperationWorker, + PanelRuntime, + PlanStatus, + TargetEvaluationWorker, + TargetStatus, +) +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + + +class EmptyServer: + pass + + +class FakeOperatorBackend: + def __init__(self) -> None: + self.cancel_calls = 0 + + def cancel(self) -> bool: + self.cancel_calls += 1 + return True + + +class FakeOperator: + def __init__(self, module: FakeOperatorBackend | None = None) -> None: + self.module = module or FakeOperatorBackend() + + def status(self) -> OperatorStatus: + return OperatorStatus(state="IDLE", error="", has_plan=False) + + def get_init_joints(self, robot_name: str) -> None: + _ = robot_name + return None + + def cancel(self) -> bool: + return self.module.cancel() + + def preview(self, *_args: object, **_kwargs: object) -> bool: + return True + + +@dataclass +class FakeStopOperationWorker(OperationWorker): + stop_calls: list[float | None] + + def __init__(self, stop_calls: list[float | None]) -> None: + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +@dataclass +class FakeStopEvaluationWorker(TargetEvaluationWorker): + stop_calls: list[float | None] + + def __init__(self, stop_calls: list[float | None]) -> None: + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +class FakeTimeoutSubmitWorker(OperationWorker): + def __init__(self, submissions: list[dict[str, float]]) -> None: + self.submissions = submissions + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + kwargs = {} + if timeout_seconds is not None: + kwargs["timeout_seconds"] = timeout_seconds + self.submissions.append(kwargs) + + +class FakeOperationSubmitWorker(OperationWorker): + def __init__(self, submissions: list[Callable[[], None]]) -> None: + self.submissions = submissions + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + self.submissions.append(operation) + + +class FakeOperationErrorWorker(OperationWorker): + def __init__(self, errors: list[Callable[[str], None]]) -> None: + self.errors = errors + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + _ = operation, timeout_seconds + if on_error is not None: + self.errors.append(on_error) + + +class FakeRestartableOperationWorker(FakeOperationSubmitWorker): + def __init__( + self, submissions: list[Callable[[], None]], stop_calls: list[float | None] + ) -> None: + super().__init__(submissions) + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +def planning_group(robot: str, name: str, joints: tuple[str, ...]) -> PlanningGroup: + return PlanningGroup( + f"{robot}/{name}", + robot, + name, + tuple(f"{robot}/{joint}" for joint in joints), + joints, + "base", + None, + ) + + +def make_gui(module: FakeOperatorBackend | None = None) -> ViserPanelGui: + module = module or FakeOperatorBackend() + return ViserPanelGui( + EmptyServer(), + PlanningSceneInfo(robots={}), + FakeOperator(module), + {}, + ViserVisualizationConfig(), + ) + + +@pytest.mark.parametrize( + ("result", "success", "collision_free", "expected"), + [ + ( + TargetEvaluationResult(True, "FEASIBLE", "", True), + True, + True, + FeasibilityStatus.FEASIBLE, + ), + (TargetEvaluationResult(False, "COLLISION", ""), False, False, FeasibilityStatus.COLLISION), + ( + TargetEvaluationResult(False, "COLLISION_AT_START", ""), + False, + False, + FeasibilityStatus.COLLISION, + ), + ( + TargetEvaluationResult(False, "COLLISION_AT_GOAL", ""), + False, + False, + FeasibilityStatus.COLLISION, + ), + ( + TargetEvaluationResult(False, "NO_SOLUTION", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + ( + TargetEvaluationResult(False, "SINGULARITY", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + ( + TargetEvaluationResult(False, "JOINT_LIMITS", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + (TargetEvaluationResult(False, "TIMEOUT", ""), False, False, FeasibilityStatus.IK_FAILED), + ( + TargetEvaluationResult(False, "IK_SUCCEEDED", ""), + False, + False, + FeasibilityStatus.INVALID, + ), + ], +) +def test_gui_feasibility_status_uses_exact_status_mapping( + result: TargetEvaluationResult, + success: bool, + collision_free: bool, + expected: FeasibilityStatus, +) -> None: + gui = make_gui() + + assert gui._feasibility_status(result, success, collision_free) == expected + + +def test_group_status_composes_shared_panel_state_without_robot_dropdown() -> None: + gui = make_gui() + values: dict[str, str] = {} + gui.state.selected_group_ids = ("left/manipulator", "right/gripper") + gui.state.error = "planner unavailable" + gui.state.target_status = gui.state.target_status.FEASIBLE + gui.state.plan_state.status = gui.state.plan_state.status.FRESH + gui._stale_robot_names = lambda _group_ids: ("right",) # type: ignore[method-assign] + gui._set_handle_value = values.__setitem__ # type: ignore[method-assign] + + gui._update_status_text() + + assert "robot" not in gui._handles + assert values == { + "status": "### Status\n\n**State:** planner unavailable\n\n" + "Target: `feasible` · Plan: `fresh`\n\nState stale: `True (right)`", + "target_summary": "Feasibility: `unknown`", + } + + +def test_gui_close_uses_bounded_operation_worker_stop(monkeypatch: pytest.MonkeyPatch) -> None: + stop_timeouts: list[float | None] = [] + gui = make_gui() + gui._operation_worker.stop() + gui._worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeStopOperationWorker(stop_timeouts)) + monkeypatch.setattr(gui, "_worker", FakeStopEvaluationWorker([])) + + gui.close() + + assert stop_timeouts == [2.0] + + +def test_gui_only_preview_submits_timeout_override(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[dict[str, float]] = [] + gui = make_gui() + gui.config = ViserVisualizationConfig(preview_request_timeout=0.25) + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeTimeoutSubmitWorker(submissions)) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.plan_state.status = PlanStatus.FRESH + gui._submit_preview() + + assert submissions == [{"timeout_seconds": 0.25}] + + +def test_gui_preview_enters_previewing_before_worker_runs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + monkeypatch.setattr(gui, "refresh", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "COMPLETED" + gui.state.selected_group_ids = ("arm/manipulator",) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + assert gui.state.can_execute() is True + + gui._submit_preview() + + assert gui.state.action_status == ActionStatus.PREVIEWING + assert gui.state.can_execute() is False + assert len(submissions) == 1 + + +def test_gui_selection_change_clears_invalidated_preview( + monkeypatch: pytest.MonkeyPatch, +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + monkeypatch.setattr(gui, "refresh", lambda: None) + groups = [ + planning_group("arm", "manipulator", ("j1",)), + planning_group("arm", "gripper", ("j2",)), + ] + monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) + monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "COMPLETED" + gui.state.selected_group_ids = (groups[0].id,) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + gui._submit_preview() + gui._toggle_group_selected(groups[1].id) + submissions[0]() + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.last_result == "preview=False" + + +def test_gui_selection_change_ignores_invalidated_preview_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + errors: list[Callable[[str], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationErrorWorker(errors)) + monkeypatch.setattr(gui, "refresh", lambda: None) + groups = [ + planning_group("arm", "manipulator", ("j1",)), + planning_group("arm", "gripper", ("j2",)), + ] + monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) + monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "COMPLETED" + gui.state.selected_group_ids = (groups[0].id,) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + gui._submit_preview() + gui._toggle_group_selected(groups[1].id) + errors[0]("preview timed out") + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.error == "" + assert gui.state.last_result == "preview=False" + + +def test_gui_cancel_bypasses_operation_worker(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[Callable[[], None]] = [] + stop_calls: list[float | None] = [] + module = FakeOperatorBackend() + gui = make_gui(module) + gui._operation_worker.stop() + monkeypatch.setattr( + gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) + ) + gui.state.action_status = ActionStatus.PREVIEWING + + gui._submit_cancel() + gui.close() + + assert submissions == [] + assert stop_calls == [0.0] + assert module.cancel_calls == 1 + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.last_result == "cancel=True" + + +def test_gui_cancelled_planning_clears_active_plan_state(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[Callable[[], None]] = [] + stop_calls: list[float | None] = [] + module = FakeOperatorBackend() + gui = make_gui(module) + gui._operation_worker.stop() + monkeypatch.setattr( + gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) + ) + stale_operation_id = gui._next_operation_id() + gui.state.action_status = ActionStatus.RUNNING + gui.state.plan_state.status = PlanStatus.PLANNING + assert gui.state.plan_state.status == PlanStatus.PLANNING + + gui._submit_cancel() + gui._finish_operation("plan_to_joints=True", operation_id=stale_operation_id) + gui.close() + + assert submissions == [] + assert module.cancel_calls == 1 + assert stop_calls == [0.0] + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.plan_state.status == PlanStatus.FAILED + assert gui.state.last_result == "cancel=True" + + +@pytest.mark.parametrize( + ("submit", "expected_error"), + [ + ("_submit_plan", "Cannot plan until target is feasible and manipulation is idle"), + ("_submit_preview", "No fresh plan to preview"), + ( + "_submit_execute", + "Cannot execute: require feasible fresh plan", + ), + ], +) +def test_gui_guard_errors_keep_action_idle( + submit: str, expected_error: str, monkeypatch: pytest.MonkeyPatch +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.selected_robot = "arm" + gui.state.action_status = ActionStatus.IDLE + + getattr(gui, submit)() + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.error == expected_error + assert submissions == [] + + +def test_gui_ignores_stale_timed_out_operation_finish() -> None: + gui = make_gui() + old_operation_id = gui._next_operation_id() + gui._set_operation_error("Operation timed out after 5.0s", old_operation_id) + gui.state.action_status = ActionStatus.FAILED + + gui._finish_operation("preview=True", operation_id=old_operation_id) + + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "Operation timed out after 5.0s" diff --git a/dimos/manipulation/visualization/viser/test_gui_status.py b/dimos/manipulation/visualization/viser/test_gui_status.py deleted file mode 100644 index f461565bdf..0000000000 --- a/dimos/manipulation/visualization/viser/test_gui_status.py +++ /dev/null @@ -1,63 +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. - -from __future__ import annotations - -import pytest - -pytest.importorskip("viser", reason="Viser optional dependency is not installed") - -from dimos.manipulation.visualization.types import TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter -from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.state import FeasibilityStatus - - -class StatusOnlyServer: - pass - - -class StatusOnlyAdapter(InProcessViserAdapter): - def __init__(self) -> None: - pass - - -@pytest.mark.parametrize( - ("result", "success", "collision_free", "expected"), - [ - ({"status": "FEASIBLE"}, True, True, FeasibilityStatus.FEASIBLE), - ({"status": "COLLISION"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "COLLISION_AT_START"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "COLLISION_AT_GOAL"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "NO_SOLUTION"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "SINGULARITY"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "JOINT_LIMITS"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "TIMEOUT"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "IK_SUCCEEDED"}, False, False, FeasibilityStatus.INVALID), - ], -) -def test_gui_feasibility_status_uses_exact_status_mapping( - result: TargetEvaluation, - success: bool, - collision_free: bool, - expected: FeasibilityStatus, -) -> None: - gui = ViserPanelGui( - StatusOnlyServer(), - StatusOnlyAdapter(), - ViserVisualizationConfig(), - ) - - assert gui._feasibility_status(result, success, collision_free) == expected diff --git a/dimos/manipulation/visualization/viser/test_operation_worker.py b/dimos/manipulation/visualization/viser/test_operation_worker.py deleted file mode 100644 index 7b7724ce2d..0000000000 --- a/dimos/manipulation/visualization/viser/test_operation_worker.py +++ /dev/null @@ -1,322 +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. - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -import threading - -import pytest - -pytest.importorskip("viser", reason="Viser optional dependency is not installed") - -from dimos.manipulation.visualization.types import RobotInfo -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter -from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.state import ( - ActionStatus, - BackendConnectionStatus, - OperationWorker, - PanelRuntime, - PlanStatus, - TargetEvaluationWorker, - TargetStatus, -) -from dimos.msgs.sensor_msgs.JointState import JointState - - -class EmptyServer: - pass - - -@dataclass -class FakeStopOperationWorker(OperationWorker): - stop_calls: list[float | None] - - def __init__(self, stop_calls: list[float | None]) -> None: - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -@dataclass -class FakeStopEvaluationWorker(TargetEvaluationWorker): - stop_calls: list[float | None] - - def __init__(self, stop_calls: list[float | None]) -> None: - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -class FakeTimeoutSubmitWorker(OperationWorker): - def __init__(self, submissions: list[dict[str, float]]) -> None: - self.submissions = submissions - - def submit( - self, - operation: Callable[[], None], - *, - timeout_seconds: float | None = None, - on_error: Callable[[str], None] | None = None, - ) -> None: - kwargs = {} - if timeout_seconds is not None: - kwargs["timeout_seconds"] = timeout_seconds - self.submissions.append(kwargs) - - -class FakeOperationSubmitWorker(OperationWorker): - def __init__(self, submissions: list[Callable[[], None]]) -> None: - self.submissions = submissions - - def submit( - self, - operation: Callable[[], None], - *, - timeout_seconds: float | None = None, - on_error: Callable[[str], None] | None = None, - ) -> None: - self.submissions.append(operation) - - -class FakeRestartableOperationWorker(FakeOperationSubmitWorker): - def __init__( - self, submissions: list[Callable[[], None]], stop_calls: list[float | None] - ) -> None: - super().__init__(submissions) - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -class FakeOperationAdapter(InProcessViserAdapter): - def __init__(self) -> None: - self.cancel_calls = 0 - - def list_robots(self) -> list[str]: - return [] - - def get_module_state(self) -> str: - return "IDLE" - - def get_robot_info(self, robot_name: str) -> RobotInfo | None: - return None - - def get_current_joint_state(self, robot_name: str) -> None: - return None - - def get_ee_pose(self, robot_name: str, joint_state: JointState | None = None) -> None: - return None - - def get_error(self) -> str: - return "" - - def get_robot_config(self, robot_name: str) -> None: - return None - - def is_state_stale(self, robot_name: str, max_age: float = 1.0) -> bool: - return False - - def cancel(self) -> bool: - self.cancel_calls += 1 - return True - - def plan_to_joints(self, joints: JointState, robot_name: str | None = None) -> bool: - return True - - -def test_operation_worker_uses_per_operation_timeout() -> None: - errors: list[str] = [] - worker = OperationWorker(errors.append, timeout_seconds=1.0) - worker.submit(lambda: None, timeout_seconds=0.25) - - request = worker._requests.get_nowait() - - assert worker._operation_timeout(request) == 0.25 - - -def test_operation_worker_uses_operation_error_callback_on_timeout() -> None: - default_errors: list[str] = [] - operation_errors: list[str] = [] - release = threading.Event() - finished = threading.Event() - worker = OperationWorker(default_errors.append) - - def operation() -> None: - release.wait(timeout=1.0) - finished.set() - - worker.submit( - operation, - timeout_seconds=0.001, - on_error=operation_errors.append, - ) - - worker._run_operation(worker._requests.get_nowait()) - release.set() - assert finished.wait(timeout=1.0) - - assert default_errors == [] - assert operation_errors == ["Operation timed out after 0.0s"] - - -def test_gui_close_uses_bounded_operation_worker_stop(monkeypatch: pytest.MonkeyPatch) -> None: - stop_timeouts: list[float | None] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - gui._worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeStopOperationWorker(stop_timeouts)) - monkeypatch.setattr(gui, "_worker", FakeStopEvaluationWorker([])) - - gui.close() - - assert stop_timeouts == [2.0] - - -def test_gui_only_preview_submits_timeout_override(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[dict[str, float]] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(preview_request_timeout=0.25), - ) - gui._operation_worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeTimeoutSubmitWorker(submissions)) - gui.state.runtime = PanelRuntime.RUNNING - gui.state.backend_status = BackendConnectionStatus.READY - gui.state.selected_robot = "arm" - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "IDLE" - - gui._submit_plan() - gui.state.plan_state.status = PlanStatus.FRESH - gui._submit_preview() - - assert "timeout_seconds" not in submissions[0] - assert submissions[1]["timeout_seconds"] == 0.25 - - -def test_gui_cancel_bypasses_operation_worker(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[Callable[[], None]] = [] - stop_calls: list[float | None] = [] - adapter = FakeOperationAdapter() - gui = ViserPanelGui( - EmptyServer(), - adapter, - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) - ) - gui.state.action_status = ActionStatus.PREVIEWING - - gui._submit_cancel() - gui.close() - - assert submissions == [] - assert stop_calls == [0.0] - assert adapter.cancel_calls == 1 - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.last_result == "cancel=True" - - -def test_gui_cancelled_planning_clears_active_plan_state(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[Callable[[], None]] = [] - stop_calls: list[float | None] = [] - adapter = FakeOperationAdapter() - gui = ViserPanelGui( - EmptyServer(), - adapter, - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) - ) - stale_operation_id = gui._next_operation_id() - gui.state.action_status = ActionStatus.RUNNING - gui.state.plan_state.status = PlanStatus.PLANNING - assert gui.state.plan_state.status == PlanStatus.PLANNING - - gui._submit_cancel() - gui._finish_operation("plan_to_joints=True", operation_id=stale_operation_id) - gui.close() - - assert submissions == [] - assert adapter.cancel_calls == 1 - assert stop_calls == [0.0] - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.plan_state.status == PlanStatus.FAILED - assert gui.state.last_result == "cancel=True" - - -@pytest.mark.parametrize( - ("submit", "expected_error"), - [ - ("_submit_plan", "Cannot plan until target is feasible and manipulation is idle"), - ("_submit_preview", "No fresh plan to preview"), - ( - "_submit_execute", - "Cannot execute: require feasible fresh plan and matching current joints", - ), - ], -) -def test_gui_guard_errors_keep_action_idle( - submit: str, expected_error: str, monkeypatch: pytest.MonkeyPatch -) -> None: - submissions: list[Callable[[], None]] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) - gui.state.runtime = PanelRuntime.RUNNING - gui.state.backend_status = BackendConnectionStatus.READY - gui.state.selected_robot = "arm" - gui.state.action_status = ActionStatus.IDLE - - getattr(gui, submit)() - - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.error == expected_error - assert submissions == [] - - -def test_gui_ignores_stale_timed_out_operation_finish() -> None: - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - old_operation_id = gui._next_operation_id() - gui._set_operation_error("Operation timed out after 5.0s", old_operation_id) - gui.state.action_status = ActionStatus.FAILED - - gui._finish_operation("preview=True", operation_id=old_operation_id) - - assert gui.state.action_status == ActionStatus.FAILED - assert gui.state.error == "Operation timed out after 5.0s" diff --git a/dimos/manipulation/visualization/viser/test_state.py b/dimos/manipulation/visualization/viser/test_state.py index 412c227050..a9ec709b93 100644 --- a/dimos/manipulation/visualization/viser/test_state.py +++ b/dimos/manipulation/visualization/viser/test_state.py @@ -14,21 +14,109 @@ from __future__ import annotations +from collections.abc import Callable +import threading + +from dimos.manipulation.planning.spec.models import PlanningGroupID from dimos.manipulation.visualization.viser.state import ( + ActionStatus, BackendConnectionStatus, + OperationWorker, PanelRuntime, PanelState, + PlanStatus, + TargetEvaluationWorker, TargetStatus, ) -def test_panel_can_plan_from_fault_after_planning_failure() -> None: +def test_panel_cannot_plan_from_fault_without_explicit_reset() -> None: state = PanelState( selected_robot="arm", + selected_group_ids=(PlanningGroupID("arm/manipulator"),), runtime=PanelRuntime.RUNNING, backend_status=BackendConnectionStatus.READY, target_status=TargetStatus.FEASIBLE, manipulation_state="FAULT", ) - assert state.can_plan() is True + assert state.can_plan() is False + + +def test_panel_cannot_plan_without_a_selected_group() -> None: + state = PanelState( + runtime=PanelRuntime.RUNNING, + backend_status=BackendConnectionStatus.READY, + target_status=TargetStatus.FEASIBLE, + manipulation_state="IDLE", + ) + + assert state.can_plan() is False + + +def test_sequence_change_marks_a_fresh_plan_stale() -> None: + state = PanelState(plan_state=PanelState().plan_state) + state.plan_state.status = PlanStatus.FRESH + + state.next_sequence_id() + + assert state.plan_state.status == PlanStatus.STALE + assert state.target_status == TargetStatus.CHECKING + + +def test_selection_epoch_change_resets_plan_and_invalidates_sequence() -> None: + state = PanelState(selected_group_ids=(PlanningGroupID("arm/manipulator"),)) + state.plan_state.status = PlanStatus.FRESH + + assert state.advance_selection_epoch() == 1 + + assert state.latest_sequence_id == 1 + assert state.plan_state.status == PlanStatus.NONE + + +def test_cancel_is_available_for_running_preview_and_execute_actions() -> None: + state = PanelState() + + for action in (ActionStatus.RUNNING, ActionStatus.PREVIEWING, ActionStatus.EXECUTING): + state.action_status = action + assert state.can_cancel() is True + + +def test_operation_worker_uses_per_operation_timeout() -> None: + errors: list[str] = [] + worker = OperationWorker(errors.append, timeout_seconds=1.0) + worker.submit(lambda: None, timeout_seconds=0.25) + + request = worker._requests.get_nowait() + + assert worker._operation_timeout(request) == 0.25 + + +def test_operation_worker_uses_operation_error_callback_on_timeout() -> None: + default_errors: list[str] = [] + operation_errors: list[str] = [] + release = threading.Event() + finished = threading.Event() + worker = OperationWorker(default_errors.append) + + def operation() -> None: + release.wait(timeout=1.0) + finished.set() + + worker.submit( + operation, + timeout_seconds=0.001, + on_error=operation_errors.append, + ) + + worker._run_operation(worker._requests.get_nowait()) + release.set() + assert finished.wait(timeout=1.0) + + assert default_errors == [] + assert operation_errors == ["Operation timed out after 0.0s"] + + +class FakeTargetEvaluationWorker(TargetEvaluationWorker): + def __init__(self, calls: list[Callable[[], None]]) -> None: + self.calls = calls diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 09976b05f8..6a1d85e432 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -1,4 +1,4 @@ -# Copyright 2025-2026 Dimensional Inc. +# 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. @@ -12,1541 +12,1217 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Hermetic contract tests for the group-aware Viser manipulation panel.""" + from __future__ import annotations from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass -import sys +from pathlib import Path import threading -from types import ModuleType, SimpleNamespace, TracebackType +from types import SimpleNamespace -import numpy as np import pytest pytest.importorskip("viser", reason="Viser optional dependency is not installed") -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningSceneInfo +from dimos.manipulation.visualization.operator import TargetEvaluationResult from dimos.manipulation.visualization.viser.animation import ( - PreviewAnimator, - interpolate_joint_path, - sampled_joint_path_frames, + GroupPreviewAnimation, + PreviewFrame, + PreviewTrack, + scaled_frame_delays, ) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.scene import ViserManipulationScene +from dimos.manipulation.visualization.viser.gui import ( + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, + ViserPanelGui, + group_display_name, +) +from dimos.manipulation.visualization.viser.scene import ( + GOAL_ROBOT_FEASIBLE_COLOR, + GOAL_ROBOT_INFEASIBLE_COLOR, + TARGET_CONTROL_FEASIBLE_COLOR, + TARGET_CONTROL_INFEASIBLE_COLOR, + ViserManipulationScene, +) from dimos.manipulation.visualization.viser.state import ( - ActionStatus, - FeasibilityStatus, - OperationWorker, PanelPlanState, PlanStatus, TargetEvaluationRequest, - TargetEvaluationWorker, TargetStatus, ) -from dimos.manipulation.visualization.viser.theme import _dimos_logo_data_url, apply_dimos_theme +from dimos.manipulation.visualization.viser.theme import apply_dimos_theme +from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState - -GuiCallback = Callable[[SimpleNamespace], None] -ThemeValue = str | bool | tuple[int, int, int] | dict[str, str | dict[str, str]] | None -RobotConfigOverride = str | list[str] | list[float] | None - - -@dataclass -class RobotConfigStub: - name: str = "arm" - joint_names: list[str] | None = None - end_effector_link: str = "ee_link" - base_link: str = "base_link" - home_joints: list[float] | None = None - joint_limits_lower: list[float] | None = None - joint_limits_upper: list[float] | None = None - - def __post_init__(self) -> None: - if self.joint_names is None: - self.joint_names = ["j1", "j2"] - - -@dataclass -class SceneRobotConfigStub: - name: str = "arm" - model_path: str = "/tmp/arm.urdf" - package_paths: dict[str, str] | None = None - xacro_args: dict[str, str] | None = None - auto_convert_meshes: bool = False - joint_names: list[str] | None = None - - def __post_init__(self) -> None: - if self.package_paths is None: - self.package_paths = {} - if self.xacro_args is None: - self.xacro_args = {} - if self.joint_names is None: - self.joint_names = ["joint1"] - - -@dataclass -class NamedState: - name: str +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint @dataclass -class GuiMarkdownHandle: - value: str +class Handle: + label: str = "" + value: object = None + options: list[str] | None = None + disabled: bool = False + color: tuple[int, int, int] | None = None + min: float = 0.0 + max: float = 0.0 + step: float = 0.0 + visible: bool = True + callback: Callable[[object], None] | None = None removed: bool = False - def remove(self) -> None: - self.removed = True + def on_update(self, callback: Callable[[object], None]) -> None: + self.callback = callback - -@dataclass -class GuiDropdownHandle: - label: str - options: list[str] - value: str - update_callback: GuiCallback | None = None - removed: bool = False - - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def on_click(self, callback: Callable[[object], None]) -> None: + self.callback = callback def remove(self) -> None: self.removed = True -@dataclass -class GuiButtonHandle: - label: str - disabled: bool = False - click_callback: GuiCallback | None = None - removed: bool = False - - def on_click(self, callback: GuiCallback) -> None: - self.click_callback = callback +class Folder(Handle): + def __init__(self, label: str, **kwargs: bool) -> None: + super().__init__(label=label) + self.kwargs = kwargs - def remove(self) -> None: - self.removed = True + def __enter__(self) -> Folder: + return self + def __exit__(self, *_: object) -> bool: + return False -@dataclass -class GuiCheckboxHandle: - label: str - value: bool - update_callback: GuiCallback | None = None - removed: bool = False - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback +class Gui: + def __init__(self) -> None: + self.folders: list[Folder] = [] + self.buttons: list[Handle] = [] + self.dropdowns: list[Handle] = [] + self.sliders: list[Handle] = [] + self.markdown: list[Handle] = [] + self.theme_kwargs: dict[str, object] | None = None + + def add_folder(self, label: str, **kwargs: bool) -> Folder: + folder = Folder(label, **kwargs) + self.folders.append(folder) + return folder + + def add_markdown(self, value: str) -> Handle: + handle = Handle(value=value) + self.markdown.append(handle) + return handle - def remove(self) -> None: - self.removed = True + def add_button(self, label: str, **kwargs: object) -> Handle: + color = kwargs.get("color") + handle = Handle( + label=label, + disabled=bool(kwargs.get("disabled", False)), + color=color if isinstance(color, tuple) else None, + ) + self.buttons.append(handle) + return handle + def add_dropdown(self, label: str, *, options: Sequence[str], initial_value: str) -> Handle: + handle = Handle(label=label, options=list(options), value=initial_value) + self.dropdowns.append(handle) + return handle -@dataclass -class GuiSliderHandle: - label: str - min: float - max: float - step: float - value: float - removed: bool = False - update_callback: GuiCallback | None = None + def add_checkbox(self, label: str, *, initial_value: bool) -> Handle: + return Handle(label=label, value=initial_value) - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def add_slider(self, label: str, **kwargs: float) -> Handle: + handle = Handle(label=label, value=kwargs["initial_value"]) + handle.min, handle.max, handle.step = kwargs["min"], kwargs["max"], kwargs["step"] + self.sliders.append(handle) + return handle - def remove(self) -> None: - self.removed = True + def configure_theme(self, **kwargs: object) -> None: + self.theme_kwargs = kwargs -class FakeHandle: +class Server: def __init__(self) -> None: - self.visible: object | None = None - self.removed = False - self.name = "" - self.kwargs: dict[str, float | bool] = {} + self.gui = Gui() + self.scene = SimpleNamespace() - def remove(self) -> None: - self.removed = True +@dataclass +class Config: + name: str + joint_names: list[str] + joint_limits_lower: list[float] + joint_limits_upper: list[float] + home_joints: list[float] | None + base_link: str = "base" + end_effector_link: str = "tool" + model_path: Path | str = "robot.urdf" + package_paths: dict[str, str] | None = None + xacro_args: dict[str, str] | None = None + auto_convert_meshes: bool = False + max_velocity: float = 1.0 + max_acceleration: float = 1.0 + joint_name_mapping: dict[str, str] | None = None + coordinator_task_name: str | None = None + pre_grasp_offset: float = 0.0 -class FakeUrdf: - def __init__(self, names: tuple[str, ...]) -> None: - self._urdf = SimpleNamespace(actuated_joint_names=names) - self._meshes = [] - self.cfg = None - self.removed = False + def __post_init__(self) -> None: + if isinstance(self.model_path, str): + self.model_path = Path(self.model_path) + + +def group(robot: str, name: str, joints: tuple[str, ...], *, pose: bool = False) -> PlanningGroup: + return PlanningGroup( + f"{robot}/{name}", + robot, + name, + tuple(f"{robot}/{joint}" for joint in joints), + joints, + "base", + "tool" if pose else None, + ) - def update_cfg(self, cfg: Sequence[float]) -> None: - self.cfg = list(cfg) - def remove(self) -> None: - self.removed = True +class Module: + def __init__(self, groups: list[PlanningGroup], states: dict[str, JointState]) -> None: + self.groups = groups + self.states = states + robots = {item.robot_name for item in groups} + self.configs = { + robot_name: Config(robot_name, ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + for robot_name in robots + } + self.plans: list[tuple[tuple[str, ...], dict[str, JointState]]] = [] + self.executions = 0 + self.cancelled = 0 + self.cleared = 0 + self.last_plan: GeneratedPlan | None = None + + def make_plan(self, group_ids: tuple[str, ...]) -> GeneratedPlan: + names = [ + name for group in self.groups for name in group.joint_names if group.id in group_ids + ] + if not names: + names = ["robot/j1", "robot/j2"] + plan = GeneratedPlan( + group_ids=group_ids, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint(0.0, [0.0] * len(names)), + TrajectoryPoint(1.0, [1.0] * len(names)), + ], + ), + path=[JointState({"name": names, "position": [0.0] * len(names)})], + status=PlanningStatus.SUCCESS, + ) + self.last_plan = plan + return plan + def list_robots(self) -> list[str]: + return list(self.configs) -class FakeJointState(JointState): - def __init__( - self, - name: Sequence[str], - position: Sequence[float] | None = None, - velocity: Sequence[float] | None = None, - effort: Sequence[float] | None = None, - ) -> None: - self.ts = 0.0 - self.frame_id = "" - self.name = list(name) - self.position = list(position or []) - self.velocity = list(velocity or []) - self.effort = list(effort or []) + def list_planning_groups(self) -> list[PlanningGroup]: + return self.groups + def robot_items(self) -> list[tuple[str, str, Config]]: + return [(name, f"id-{name}", config) for name, config in self.configs.items()] -class FakeServer: - def __init__(self) -> None: - self.scene = SimpleNamespace() - self.scene.add_transform_controls = self.add_transform_controls + def robot_id_for_name(self, name: str) -> str: + return f"id-{name}" - def add_transform_controls(self, path: str, *, scale: float) -> FakeTransformHandle: - handle = FakeTransformHandle() - handle.path = path - handle.scale = scale - return handle + def get_robot_config(self, name: str) -> Config: + return self.configs[name] + def get_init_joints(self, name: str) -> JointState: + return JointState({"name": self.configs[name].joint_names, "position": [-0.5, -1.0]}) -class FakeGridServer(FakeServer): - def __init__(self) -> None: - super().__init__() - self.grids = [] - self.scene.add_grid = self.add_grid - - def add_grid(self, name: str, **kwargs: float | bool) -> FakeHandle: - handle = FakeHandle() - handle.name = name - handle.kwargs = kwargs - handle.visible = kwargs.get("visible") - self.grids.append(handle) - return handle + def get_state(self) -> str: + return "IDLE" + def get_error(self) -> str: + return "" -class FakeTransformHandle(FakeHandle): - def __init__(self) -> None: - super().__init__() - self.position = (0.0, 0.0, 0.0) - self.wxyz = (1.0, 0.0, 0.0, 0.0) - self.color = None - self.material_color = None - self.update_callback = None - self.path = "" - self.scale = 0.0 + def reset(self) -> SimpleNamespace: + return SimpleNamespace(is_success=lambda: True) - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def plan_to_joint_targets(self, targets: dict[str, JointState]) -> bool: + self.plans.append((tuple(targets), targets)) + return True + def preview_plan(self) -> bool: + return True -class FakeTransformServer(FakeServer): - def __init__(self) -> None: - super().__init__() - self.transform_controls = [] - self.scene.add_transform_controls = self.add_transform_controls - - def add_transform_controls(self, path: str, *, scale: float) -> FakeTransformHandle: - handle = FakeTransformHandle() - handle.path = path - handle.scale = scale - self.transform_controls.append(handle) - return handle + def execute(self) -> bool: + self.executions += 1 + return True + def cancel(self) -> bool: + self.cancelled += 1 + return True -class FakeFolder: - def __init__(self, label: str, kwargs: dict[str, bool]) -> None: - self.label = label - self.kwargs = kwargs - self.entered = False - self.exited = False - self.removed = False + def clear_planned_path(self) -> bool: + self.cleared += 1 + return True - def __enter__(self) -> FakeFolder: - self.entered = True - return self - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - traceback: TracebackType | None, - ) -> bool: - self.exited = True - return False +class Monitor: + def __init__(self, module: Module) -> None: + self.module = module + self.invalid: set[str] = set() + self.stale: set[str] = set() + self.poses: dict[str, Pose] = {} - def remove(self) -> None: - self.removed = True + def get_current_joint_state(self, robot_id: str) -> JointState: + return JointState(self.module.states[robot_id.removeprefix("id-")]) + def is_state_stale(self, robot_id: str, max_age: float = 1.0) -> bool: + return robot_id in self.stale -class FakeGuiServer: - def __init__(self) -> None: - self.theme_kwargs: dict[str, ThemeValue] | None = None - self.folders = [] - self.gui = SimpleNamespace( - add_markdown=lambda value: GuiMarkdownHandle(value=value), - add_dropdown=self.add_dropdown, - add_button=self.add_button, - add_checkbox=self.add_checkbox, - add_slider=self.add_slider, - add_folder=self.add_folder, - configure_theme=self.configure_theme, - ) - self.buttons: dict[str, GuiButtonHandle] = {} - self.checkboxes: dict[str, GuiCheckboxHandle] = {} - self.sliders: list[GuiSliderHandle] = [] + def is_state_valid(self, robot_id: str, _state: JointState) -> bool: + return robot_id not in self.invalid - def configure_theme(self, **kwargs: ThemeValue) -> None: - self.theme_kwargs = kwargs + def get_group_ee_pose(self, group_id: str, _state: JointState | None = None) -> Pose: + return self.poses.get( + group_id, + Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}), + ) - def add_folder(self, label: str, **kwargs: bool) -> FakeFolder: - handle = FakeFolder(label, kwargs) - self.folders.append(handle) - return handle - def add_dropdown( - self, label: str, *, options: Sequence[str], initial_value: str - ) -> GuiDropdownHandle: - handle = GuiDropdownHandle(label=label, options=list(options), value=initial_value) - return handle +class Operator: + def __init__(self, module: Module, monitor: Monitor) -> None: + self.module = module + self.monitor = monitor - def add_button(self, label: str, *, disabled: bool = False) -> GuiButtonHandle: - handle = GuiButtonHandle(label=label, disabled=disabled) - self.buttons[label] = handle - return handle + def status(self) -> SimpleNamespace: + return SimpleNamespace( + state=self.module.get_state(), + error=self.module.get_error(), + has_plan=True, + ) - def add_checkbox(self, label: str, *, initial_value: bool) -> GuiCheckboxHandle: - handle = GuiCheckboxHandle(label=label, value=initial_value) - self.checkboxes[label] = handle - return handle + def get_init_joints(self, robot_name: str) -> JointState | None: + return self.module.get_init_joints(robot_name) - def add_slider( - self, - label: str, - *, - min: float, - max: float, - step: float, - initial_value: float, - ) -> GuiSliderHandle: - handle = GuiSliderHandle(label=label, min=min, max=max, step=step, value=initial_value) - self.sliders.append(handle) - return handle + def evaluate_joint_target(self, request: object) -> TargetEvaluationResult: + target = request.target # type: ignore[attr-defined] + group_ids = request.group_ids # type: ignore[attr-defined] + diagnostics = { + group_id: "Target is collision-free for this robot" for group_id in group_ids + } + poses = {group_id: self.monitor.get_group_ee_pose(group_id) for group_id in group_ids} + return TargetEvaluationResult( + True, + "FEASIBLE", + "Target is collision-free for each robot", + True, + tuple(group_ids), + target, + diagnostics, + poses, + ) + def evaluate_pose_target(self, request: object) -> TargetEvaluationResult: + group_ids = tuple( + dict.fromkeys((*request.pose_targets.keys(), *request.auxiliary_group_ids)) + ) # type: ignore[attr-defined] + js = JointState( + { + "name": [ + name + for group in self.module.groups + for name in group.joint_names + if group.id in group_ids + ], + "position": [ + 0.7 + for group in self.module.groups + for _ in group.joint_names + if group.id in group_ids + ], + } + ) + return TargetEvaluationResult( + True, + "FEASIBLE", + "ok", + True, + group_ids, + js, + {}, + {group_id: self.monitor.get_group_ee_pose(group_id) for group_id in group_ids}, + ) -def make_robot_config(**overrides: RobotConfigOverride) -> RobotConfigStub: - """Build a faithful RobotModelConfig stand-in with the fields the panel reads.""" - config = RobotConfigStub() - for name, value in overrides.items(): - setattr(config, name, value) - return config + def plan_to_joints(self, request: object) -> GeneratedPlan: + self.module.plan_to_joint_targets( + {group_id: JointState({"name": [], "position": []}) for group_id in request.group_ids} + ) # type: ignore[attr-defined] + return self.module.make_plan(tuple(request.group_ids)) # type: ignore[attr-defined] + def plan_to_pose(self, request: object) -> GeneratedPlan: + return self.module.make_plan(tuple(request.pose_targets)) # type: ignore[attr-defined] -class FakeManipulationModule(SimpleNamespace): - """Public ManipulationModule surface used by the in-process Viser adapter tests.""" + def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: + _ = plan, duration + return self.module.preview_plan() - def list_robots(self) -> list[str]: - return list(getattr(self, "_robots", {}).keys()) + def execute(self, plan: GeneratedPlan) -> bool: + _ = plan + return self.module.execute() - def robot_items(self) -> list[tuple[str, str, RobotConfigStub | SimpleNamespace]]: - return [ - (name, robot_id, config) - for name, (robot_id, config, _) in getattr(self, "_robots", {}).items() - ] + def cancel(self) -> bool: + return self.module.cancel() - def robot_id_for_name(self, robot_name: str) -> str | None: - entry = getattr(self, "_robots", {}).get(robot_name) - return entry[0] if entry is not None else None - - def robot_name_for_id(self, robot_id: str) -> str | None: - for robot_name, (candidate_id, _, _) in getattr(self, "_robots", {}).items(): - if candidate_id == robot_id: - return robot_name - return None - - def get_robot_config(self, robot_name: str) -> RobotConfigStub | SimpleNamespace | None: - entry = getattr(self, "_robots", {}).get(robot_name) - return entry[1] if entry is not None else None - - def get_robot_info(self, robot_name: str) -> RobotInfo | None: - config = self.get_robot_config(robot_name) - if config is None: - return None - init = self.get_init_joints(robot_name) - home_joints = config.home_joints if hasattr(config, "home_joints") else None - return { - "name": config.name, - "world_robot_id": self.robot_id_for_name(robot_name) or robot_name, - "joint_names": list(config.joint_names), - "end_effector_link": config.end_effector_link, - "base_link": config.base_link, - "max_velocity": 1.0, - "max_acceleration": 1.0, - "has_joint_name_mapping": False, - "coordinator_task_name": None, - "home_joints": list(home_joints) if home_joints is not None else None, - "pre_grasp_offset": 0.0, - "init_joints": list(init.position) if init is not None else None, - } + def clear_plan(self) -> bool: + return self.module.clear_planned_path() - def get_init_joints(self, robot_name: str) -> JointState | None: - return getattr(self, "_init_joints", {}).get(robot_name) + def reset(self) -> bool: + result = self.module.reset() + return result.is_success() - def get_state(self) -> str: - state = getattr(self, "_state", "IDLE") - return str(getattr(state, "name", state)) - def get_error(self) -> str: - return str(getattr(self, "_error_message", "")) - - def evaluate_joint_target(self, joints: JointState | None, robot_name: str) -> TargetEvaluation: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None or joints is None: - return {"success": False, "status": "NO_ROBOT", "joint_state": None} - world_monitor = getattr(self, "_world_monitor", None) - if world_monitor is None: - return {"success": False, "status": "UNAVAILABLE", "joint_state": None} - collision_free = world_monitor.is_state_valid(robot_id, joints) - return { - "success": True, - "status": "FEASIBLE" if collision_free else "COLLISION", - "message": "Target is collision-free" if collision_free else "Target is in collision", - "collision_free": collision_free, - "ee_pose": world_monitor.get_ee_pose(robot_id, joints), - "joint_state": joints, - } +def session_inputs(module: Module) -> tuple[PlanningSceneInfo, Operator, dict[str, JointState]]: + monitor = Monitor(module) + robots = {f"id-{name}": config for name, config in module.configs.items()} + current = {f"id-{name}": JointState(state) for name, state in module.states.items()} + return ( + PlanningSceneInfo(robots=robots, planning_groups=tuple(module.groups)), + Operator(module, monitor), + current, + ) - def evaluate_pose_target(self, _pose: Pose, _robot_name: str) -> TargetEvaluation: - return { - "success": False, - "joint_state": None, - "status": "UNAVAILABLE", - "message": "No fake pose IK", - "collision_free": False, - } - -def make_adapter_with_robot() -> InProcessViserAdapter: - current = FakeJointState(["j1", "j2"], position=[0.3, 0.4]) - config = make_robot_config( - name="arm", - joint_names=["j1", "j2"], - joint_limits_lower=[-1.0, -2.0], - joint_limits_upper=[1.0, 2.0], - home_joints=[0.0, 0.0], - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda _robot_id: current, - is_state_stale=lambda _robot_id, max_age=1.0: False, - is_state_valid=lambda _robot_id, _joint_state: True, - get_ee_pose=lambda _robot_id, joint_state=None: None, - ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - _init_joints={"arm": FakeJointState(["j1", "j2"], position=[0.1, 0.2])}, - _state=NamedState(name="IDLE"), - _error_message="", - _world_monitor=world_monitor, - ) - return InProcessViserAdapter( - world_monitor=world_monitor, - manipulation_module=module, - ) +def scene_gui(module: Module, server: Server, scene: ViserManipulationScene) -> ViserPanelGui: + scene_info, operator, current = session_inputs(module) + return ViserPanelGui(server, scene_info, operator, current, ViserVisualizationConfig(), scene) @pytest.fixture -def make_panel() -> Iterator[Callable[..., ViserPanelGui]]: - """Build and start a ViserPanelGui, closing it (and its worker threads) on teardown.""" +def panel() -> Iterator[ + Callable[[list[PlanningGroup], dict[str, JointState]], tuple[ViserPanelGui, Module, Server]] +]: panels: list[ViserPanelGui] = [] - def _make( - server: FakeGuiServer | FakeServer, - adapter: InProcessViserAdapter, - config: ViserVisualizationConfig | None = None, - scene: ViserManipulationScene | None = None, - ) -> ViserPanelGui: + def make( + groups: list[PlanningGroup], states: dict[str, JointState] + ) -> tuple[ViserPanelGui, Module, Server]: + module = Module(groups, states) + server = Server() + scene_info, operator, current = session_inputs(module) gui = ViserPanelGui( - server, adapter, config or ViserVisualizationConfig(panel_enabled=True), scene + server, scene_info, operator, current, ViserVisualizationConfig(panel_enabled=True) ) gui.start() panels.append(gui) - return gui + return gui, module, server - yield _make + yield make for gui in panels: gui.close() -def test_viser_config_enables_panel_by_default() -> None: - assert ViserVisualizationConfig().panel_enabled is True +def states(*robots: str) -> dict[str, JointState]: + return {robot: JointState({"name": ["j1", "j2"], "position": [0.1, 0.2]}) for robot in robots} -def test_gui_builds_controls_in_manipulation_panel_folder( - make_panel: Callable[..., ViserPanelGui], -) -> None: - server = FakeGuiServer() - adapter = make_adapter_with_robot() - gui = make_panel(server, adapter, ViserVisualizationConfig()) - assert server.folders - assert server.folders[0].label == "Manipulation Panel" - assert server.folders[0].kwargs == {"expand_by_default": True} - assert "status" in gui._handles - assert "robot" in gui._handles - assert "plan" in gui._handles - assert gui._operation_worker._timeout_seconds is None - - -def test_gui_scene_grid_checkbox_toggles_reference_grid( - make_panel: Callable[..., ViserPanelGui], +def test_panel_contract_group_order_defaults_and_controls( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - grid_server = FakeGridServer() - scene = ViserManipulationScene( - grid_server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 - ) - server = FakeGuiServer() - adapter = make_adapter_with_robot() - make_panel(server, adapter, ViserVisualizationConfig(), scene) - assert grid_server.grids - assert server.checkboxes["Scene grid"].value is True - server.checkboxes["Scene grid"].update_callback( - SimpleNamespace(target=SimpleNamespace(value=False)) - ) - assert grid_server.grids[0].visible is False - server.checkboxes["Scene grid"].update_callback( - SimpleNamespace(target=SimpleNamespace(value=True)) - ) - assert grid_server.grids[0].visible is True - + pose = group("arm", "manipulator", ("j1",), pose=True) + auxiliary = group("arm", "gripper", ("j2",)) + gui, _module, server = panel([auxiliary, pose], states("arm")) -def test_gui_close_removes_handles_and_late_callbacks_are_noops( - make_panel: Callable[..., ViserPanelGui], -) -> None: - server = FakeGuiServer() - grid_server = FakeGridServer() - scene = ViserManipulationScene( - grid_server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 + assert [(folder.label, folder.kwargs) for folder in server.gui.folders] == [ + ("Manipulation Panel", {"expand_by_default": True}), + ("Joint Control", {"expand_by_default": False}), + ] + assert [button.label for button in server.gui.buttons] == [ + "arm", + "arm gripper", + "Plan", + "Preview", + "Execute", + "Cancel", + "Clear plan", + ] + assert "robot" not in gui._handles + assert ( + server.gui.markdown[1].value + == "### Planning Groups\nActive planning groups for pose goals, planning, and joint edits." ) - adapter = make_adapter_with_robot() - gui = make_panel(server, adapter, ViserVisualizationConfig(), scene) - robot_dropdown = gui._handles["robot"] - plan_button = server.buttons["Plan"] - grid = grid_server.grids[0] - handles = list(gui._handles.values()) - - gui.close() - if isinstance(robot_dropdown, GuiDropdownHandle) and robot_dropdown.update_callback is not None: - robot_dropdown.update_callback(SimpleNamespace(target=SimpleNamespace(value="arm"))) - if plan_button.click_callback is not None: - plan_button.click_callback(SimpleNamespace()) - gui._set_scene_grid_visible(False) - - assert all(getattr(handle, "removed", False) for handle in handles) - assert gui._handles == {} - assert grid.visible is True + assert [button.color for button in server.gui.buttons[:2]] == [ + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, + ] + assert gui.state.selected_group_ids == ("arm/manipulator",) + assert server.gui.dropdowns[0].options == ["Select preset...", "Init", "Current", "Home"] + assert [ + (slider.label, slider.min, slider.max, slider.value) for slider in server.gui.sliders + ] == [("arm/manipulator/j1", -1.0, 1.0, 0.1)] + server.gui.buttons[1].callback(SimpleNamespace()) + assert gui.state.selected_group_ids == ("arm/manipulator", "arm/gripper") + assert [slider.label for slider in server.gui.sliders if not slider.removed] == [ + "arm/manipulator/j1", + "arm/gripper/j2", + ] -def test_gui_ignores_target_evaluation_after_close( - make_panel: Callable[..., ViserPanelGui], +def test_gui_target_ghost_states_use_exact_group_names() -> None: + left, right = group("left", "manipulator", ("j1",)), group("right", "manipulator", ("j1",)) + module = Module([left, right], states("left", "right")) + scene_info, operator, current = session_inputs(module) + gui = ViserPanelGui(Server(), scene_info, operator, current, ViserVisualizationConfig()) + gui.state.selected_group_ids = (left.id, right.id) + targets = { + left.id: JointState({"name": ["left/j1"], "position": [0.7]}), + right.id: JointState({"name": ["right/j1"], "position": [0.8]}), + } + ghost_states = gui._target_ghost_states(targets) + assert ghost_states["left"].position == [0.7, 0.2] + assert ghost_states["right"].position == [0.8, 0.2] + assert gui.evaluate_joint_target_set((left.id, right.id), targets).status == "FEASIBLE" + + +def test_target_callbacks_require_current_sequence_and_selection_epoch( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui.state.selected_robot = "arm" - sequence_id = gui.state.next_sequence_id() - request = TargetEvaluationRequest( - sequence_id=sequence_id, - source="joints", - robot_name="arm", - joints=FakeJointState(["j1", "j2"], position=[0.1, 0.2]), + first, second = ( + group("arm", "manipulator", ("j1",), pose=True), + group("arm", "gripper", ("j2",)), ) - gui.close() - - gui._apply_target_evaluation_result( - request, - { - "success": True, - "collision_free": True, - "status": "FEASIBLE", - "joint_state": FakeJointState(["j1", "j2"], position=[0.8, 0.9]), - }, + gui, _module, _server = panel([first, second], states("arm")) + request = TargetEvaluationRequest( + 1, "joints", selection_epoch=gui.state.selection_epoch, group_ids=(first.id,) ) - + gui.state.latest_sequence_id = 2 + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) + assert gui.state.target_status == TargetStatus.EMPTY + gui.state.latest_sequence_id = 1 + gui.state.advance_selection_epoch() + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) assert gui.state.target_status == TargetStatus.CHECKING - assert gui.state.joint_target is None - - -def test_dimos_theme_configures_supported_viser_chrome() -> None: - server = FakeGuiServer() - assert apply_dimos_theme(server) is True - assert server.theme_kwargs is not None - assert server.theme_kwargs["brand_color"] == (22, 130, 163) - assert server.theme_kwargs["dark_mode"] is True - assert server.theme_kwargs["show_logo"] is False - assert server.theme_kwargs["show_share_button"] is False - assert server.theme_kwargs["control_layout"] == "collapsible" - assert server.theme_kwargs["control_width"] == "medium" - - -def test_dimos_theme_configures_titlebar_when_supported(monkeypatch: pytest.MonkeyPatch) -> None: - fake_viser = ModuleType("viser") - fake_theme = ModuleType("viser.theme") - fake_theme.TitlebarImage = lambda **kwargs: kwargs - fake_theme.TitlebarButton = lambda **kwargs: kwargs - fake_theme.TitlebarConfig = lambda **kwargs: kwargs - monkeypatch.setitem(sys.modules, "viser", fake_viser) - monkeypatch.setitem(sys.modules, "viser.theme", fake_theme) - server = FakeGuiServer() - - assert apply_dimos_theme(server) is True - assert server.theme_kwargs is not None - titlebar_content = server.theme_kwargs["titlebar_content"] - assert isinstance(titlebar_content, dict) - image = titlebar_content["image"] - assert isinstance(image, dict) - assert image["image_alt"] == "Dimensional" - assert image["image_url_light"].startswith("data:image/svg+xml;base64,") - -def test_dimos_logo_asset_loads_as_data_url() -> None: - logo_url = _dimos_logo_data_url() - - assert logo_url is not None - assert logo_url.startswith("data:image/svg+xml;base64,") +def test_plan_target_sequence_invalidation_and_unfiltered_all_robot_execute( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], monkeypatch: pytest.MonkeyPatch +) -> None: + left, right = ( + group("left", "manipulator", ("j1",), pose=True), + group("right", "manipulator", ("j1",), pose=True), + ) + gui, module, _server = panel([left, right], states("left", "right")) + gui._toggle_group_selected(right.id) + gui.state.target_status = TargetStatus.FEASIBLE + monkeypatch.setattr( + gui, + "_operation_worker", + SimpleNamespace(submit=lambda operation, **_: operation(), stop=lambda **_: None), + ) + gui._submit_plan() + assert module.plans[-1][0] == (left.id, right.id) + assert gui.state.plan_state.group_ids == (left.id, right.id) + gui.state.next_sequence_id() + assert gui.state.plan_state.status == PlanStatus.STALE + gui.state.plan_state = PanelPlanState( + status=PlanStatus.FRESH, + group_ids=(left.id, right.id), + target_sequence_id=gui.state.latest_sequence_id, + plan=module.last_plan, + ) + gui.state.target_status = TargetStatus.FEASIBLE + gui._submit_execute() + assert module.executions == 1 -def test_dimos_theme_is_non_blocking_when_theme_api_fails() -> None: - class BrokenGui: - @staticmethod - def configure_theme(**_kwargs: ThemeValue) -> None: - raise TypeError("theme API changed") +def test_initialization_waits_for_complete_fresh_telemetry( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel( + [selected], {"arm": JointState({"name": ["j1"], "position": [0.4]})} + ) - server = SimpleNamespace(gui=BrokenGui()) + assert selected.id not in gui.state.group_joint_targets + module.states["arm"] = JointState({"name": ["j1", "j2"], "position": [0.4, 0.5]}) + gui.refresh() + assert selected.id not in gui.state.group_joint_targets - assert apply_dimos_theme(server) is False + gui.current_states["id-arm"] = module.states["arm"] + gui.refresh() + assert gui.state.group_joint_targets[selected.id].position == [0.4, 0.5] -def test_dimos_theme_retries_without_titlebar_when_titlebar_content_fails( - monkeypatch: pytest.MonkeyPatch, +def test_incomplete_preset_preserves_existing_group_targets( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - fake_viser = ModuleType("viser") - fake_theme = ModuleType("viser.theme") - fake_theme.TitlebarImage = lambda **kwargs: kwargs - fake_theme.TitlebarButton = lambda **kwargs: kwargs - fake_theme.TitlebarConfig = lambda **kwargs: kwargs - monkeypatch.setitem(sys.modules, "viser", fake_viser) - monkeypatch.setitem(sys.modules, "viser.theme", fake_theme) - titlebar_values: list[ThemeValue] = [] - - class FallbackGui: - @staticmethod - def configure_theme(**kwargs: ThemeValue) -> None: - titlebar_values.append(kwargs["titlebar_content"]) - if kwargs["titlebar_content"] is not None: - raise TypeError("titlebar unsupported") - - server = SimpleNamespace(gui=FallbackGui()) + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel([selected], states("arm")) + before = JointState(gui.state.group_joint_targets[selected.id]) + module.get_init_joints = lambda name: JointState({"name": ["j1"], "position": [-0.5]}) - assert apply_dimos_theme(server) is True - assert titlebar_values[0] is not None - assert titlebar_values[1] is None + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id] == before + assert "missing joints" in gui.state.error -class FakeMesh: - def __init__(self) -> None: - self.visible = None - self.color = None - self.material_color = None - self.opacity = None +def test_valid_init_preset_builds_sliders_after_incomplete_initial_telemetry( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, server = panel( + [selected], {"arm": JointState({"name": ["j1"], "position": [0.4]})} + ) -class FakeViserUrdfWithMeshes: - def __init__(self, names: tuple[str, ...] = ("joint1", "joint2", "joint3")) -> None: - self._urdf = SimpleNamespace(actuated_joint_names=names) - self._meshes = [FakeMesh(), FakeMesh()] - self.cfg = None + assert gui.state.group_joint_targets == {} + assert server.gui.sliders == [] - def update_cfg(self, cfg: Sequence[float]) -> None: - self.cfg = list(cfg) + module.configs["arm"].home_joints = [-0.5, -1.0] + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id].position == [-0.5, -1.0] + assert [slider.label for slider in server.gui.sliders if not slider.removed] == [ + "arm/manipulator/j1", + "arm/manipulator/j2", + ] -def test_viser_joint_configuration_maps_names_to_urdf_order() -> None: - server = FakeServer() - urdf = FakeUrdf(("shoulder", "elbow", "wrist")) - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdf, preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - cfg = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["arm/shoulder", "elbow"], +def test_incomplete_multi_group_preset_does_not_change_any_targets( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + left, right = ( + group("left", "manipulator", ("j1",), pose=True), + group("right", "manipulator", ("j1",), pose=True), + ) + gui, module, _server = panel([left, right], states("left", "right")) + gui._toggle_group_selected(right.id) + before = { + group_id: JointState(target) for group_id, target in gui.state.group_joint_targets.items() + } + module.get_init_joints = lambda name: JointState( + {"name": ["j1"] if name == "left" else [], "position": [-0.5] if name == "left" else []} ) - scene.register_robot("robot1", cfg) - scene.set_urdf_joints(urdf, cfg.joint_names, [1.5, 2.5]) - assert urdf.cfg == [1.5, 2.5, 0.0] + gui._apply_preset("Init") -def test_scene_adds_reference_grid_when_supported() -> None: - server = FakeGridServer() - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("j1",)), preview_fps=10.0 - ) + assert gui.state.group_joint_targets == before + assert "missing joints" in gui.state.error - assert scene.has_reference_grid() is True - assert len(server.grids) == 1 - grid = server.grids[0] - assert grid.name == "/reference_grid" - assert grid.kwargs["plane"] == "xy" - assert grid.kwargs["infinite_grid"] is True - assert grid.kwargs["visible"] is True - scene.set_reference_grid_visible(False) - assert grid.visible is False - scene.set_reference_grid_visible(True) - assert grid.visible is True - - -def test_preview_visibility_only_affects_preview_ghost_and_close_removes_handles() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], - ) - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - assert all(mesh.visible is True for mesh in target._meshes) - assert all(mesh.visible is False for mesh in preview._meshes) - scene.show_preview("robot1") - assert all(mesh.visible is True for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - scene.hide_preview("robot1") - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - scene.close() - assert scene._handles == {} - assert all(mesh.visible is False for mesh in preview._meshes) - - -def test_target_ghost_is_visible_and_tracks_current_until_target_moves_it() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], - ) - scene.register_robot("robot1", config) - current = scene._urdfs["robot1:current"] - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - - assert all(mesh.visible is True for mesh in target._meshes) - assert all(mesh.visible is False for mesh in preview._meshes) - scene.update_current_robot("robot1", FakeJointState(["joint1"], position=[0.25])) - assert current.cfg == [0.25] - assert target.cfg == [0.25] - assert preview.cfg is None - - scene.set_target_joints("robot1", ["joint1"], [0.8]) - scene.update_current_robot("robot1", FakeJointState(["joint1"], position=[0.1])) - assert current.cfg == [0.1] - assert target.cfg == [0.8] - assert preview.cfg is None - - -def test_preview_animation_uses_separate_colored_ghost_and_hides_after_playback() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], +def test_cancel_clear_and_close_invalidate_operations_and_preview_generation( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], monkeypatch: pytest.MonkeyPatch +) -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + gui, module, _server = panel([selected], states("arm")) + submitted: list[Callable[[], None]] = [] + gui._operation_worker.stop() + monkeypatch.setattr( + gui, + "_operation_worker", + SimpleNamespace( + submit=lambda operation, **_: submitted.append(operation), + stop=lambda **_: None, + start=lambda: None, + ), ) - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - - assert all(mesh.color == (255, 122, 0) for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - assert all(mesh.opacity == 0.55 for mesh in preview._meshes) - - ok = scene.animate_path( - "robot1", - [ - FakeJointState(["joint1"], position=[0.0]), - FakeJointState(["joint1"], position=[1.0]), - ], - duration=0.0, + gui.state.target_status = TargetStatus.FEASIBLE + gui._submit_plan() + gui._submit_clear() + submitted[0]() + assert module.plans == [] + submitted[1]() + assert module.cleared == 1 and gui.state.plan_state.status == PlanStatus.NONE + gui.close() + status_before_callback = gui.state.target_status + gui._apply_target_evaluation_result( + TargetEvaluationRequest(0, "joints"), TargetEvaluationResult(True, "FEASIBLE", "", True) ) - - assert ok is True - assert preview.cfg == [1.0] - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) + assert gui.state.target_status is status_before_callback -def test_scene_target_helpers_handle_missing_robot_and_pose() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 - ) +class Mesh: + def __init__(self) -> None: + self.visible = False + self.color: tuple[int, int, int] | None = None + self.opacity: float | None = None - assert scene.animate_path("missing", [], duration=0.0) is False - assert scene.set_target_joints("missing", ["joint1"], [1.0]) is False - scene.set_target_pose("missing", Pose()) - handle = scene.ensure_target_controls("robot1", lambda _target: None) - scene.set_target_pose("robot1", None) - assert handle is not None - assert handle.position == (0.0, 0.0, 0.0) +class Urdf: + def __init__(self, *_: object, **__: object) -> None: + self._urdf = SimpleNamespace(actuated_joint_names=("j1", "j2")) + self._meshes = [Mesh()] + self.cfg: list[float] | None = None + def update_cfg(self, cfg: Sequence[float]) -> None: + self.cfg = list(cfg) -def test_scene_close_removes_grid_transform_and_urdf_handles() -> None: - server = FakeGridServer() - current = FakeUrdf(("joint1",)) - target = FakeUrdf(("joint1",)) - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 - ) - handle = scene.ensure_target_controls("robot1", lambda _target: None) - scene._urdfs["robot1:current"] = current - scene._urdfs["robot1:target"] = target + def remove(self) -> None: + pass - scene.close() - assert handle is not None and handle.removed is True - assert current.removed is True - assert target.removed is True - assert server.grids[0].removed is True - assert scene.has_reference_grid() is False +def test_scene_active_only_ghosts_group_gizmos_feasibility_and_shared_ticks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + updates: list[tuple[str, list[float]]] = [] + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene.set_target_active("id-arm", False) + assert scene._urdfs["id-arm:target"]._meshes[0].visible is False + scene.set_target_joints("id-arm", ["j1", "j2"], [0.8, 0.2]) + assert scene._urdfs["id-arm:target"].cfg == [0.8, 0.2] + scene.set_target_visual_state("id-arm", False) + assert scene._urdfs["id-arm:target"]._meshes[0].color == (255, 30, 30) + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", lambda _delay: None + ) + original = scene._set_preview_ghost_joints + scene._set_preview_ghost_joints = lambda robot, names, values: ( + updates.append((robot, list(values))), + original(robot, names, values), + ) # type: ignore[method-assign] + preview = GroupPreviewAnimation( + ( + PreviewTrack( + "id-arm", + ("j1", "j2"), + ( + PreviewFrame(0.0, (0.0, 0.2)), + PreviewFrame(1.0, (1.0, 0.2)), + ), + ), + ) + ) + assert scene.animate_preview(preview, 1.0) is True + assert updates[-1] == ("id-arm", [1.0, 0.2]) -def test_sampled_joint_path_frames_preserves_dense_trajectory_samples() -> None: - dense_path = [FakeJointState(["j1"], position=[float(index)]) for index in range(32)] +def test_theme_and_reference_scene_contract() -> None: + server = Server() + assert apply_dimos_theme(server) is True + assert server.gui.theme_kwargs is not None + assert server.gui.theme_kwargs["brand_color"] == (0, 153, 255) + assert server.gui.theme_kwargs["dark_mode"] is True + assert server.gui.theme_kwargs["control_layout"] == "fixed" + assert ViserVisualizationConfig().panel_enabled is True - frames = sampled_joint_path_frames(dense_path, duration=1.0, fps=30.0) - assert frames == [[float(index)] for index in range(32)] +def test_preview_selection_rejects_malformed_before_visibility() -> None: + selection = PlanningGroupSelection.from_groups((group("arm", "manipulator", ("j1",)),)) + assert selection.group_ids == ("arm/manipulator",) + # The scene transaction itself rejects missing tracks before revealing ghosts. + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + assert scene.animate_preview(GroupPreviewAnimation(()), 1.0) is False -def test_sampled_joint_path_frames_interpolates_sparse_paths() -> None: - sparse_path = [ - FakeJointState(["j1"], position=[0.0]), - FakeJointState(["j1"], position=[1.0]), +def test_group_controls_use_source_labels_and_active_colors( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + pose = group("arm", "manipulator", ("j1",), pose=True) + auxiliary = group("arm", "gripper", ("j2",)) + gui, _module, server = panel([auxiliary, pose], states("arm")) + + assert group_display_name(pose) == "arm" + assert group_display_name(auxiliary) == "arm gripper" + assert [button.label for button in server.gui.buttons[:2]] == ["arm", "arm gripper"] + assert [button.color for button in server.gui.buttons[:2]] == [ + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, ] + assert server.gui.buttons[1].callback is not None + server.gui.buttons[1].callback(SimpleNamespace()) + assert gui._handles[f"group:{auxiliary.id}"].color == ACTIVE_GROUP_COLOR - frames = sampled_joint_path_frames(sparse_path, duration=1.0, fps=4.0) - - assert frames == [[0.0], [0.25], [0.5], [0.75], [1.0]] - -def test_joint_path_frame_edge_cases_and_empty_animation() -> None: - empty_position = FakeJointState(["j1"], position=[]) - single = FakeJointState(["j1"], position=[0.7]) - start = FakeJointState(["j1"], position=[0.0]) - middle = FakeJointState(["j1"], position=[1.0]) - mismatched_final = FakeJointState(["j1", "j2"], position=[2.0, 3.0]) - set_calls: list[list[float]] = [] - sleep_calls: list[float] = [] - - assert interpolate_joint_path([empty_position], duration=1.0, fps=10.0) == [] - assert interpolate_joint_path([single], duration=1.0, fps=10.0) == [[0.7]] - assert interpolate_joint_path([start, middle, mismatched_final], duration=1.0, fps=2.0) == [ - [0.0], - [2.0, 3.0], +def test_panel_preset_defaults_and_joint_slider_limits( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + _gui, _module, server = panel([selected], states("arm")) + + assert server.gui.dropdowns[0].options == ["Select preset...", "Init", "Current", "Home"] + assert [ + (slider.label, slider.min, slider.max, slider.step, slider.value) + for slider in server.gui.sliders + ] == [ + ("arm/manipulator/j1", -1.0, 1.0, 0.001, 0.1), + ("arm/manipulator/j2", -2.0, 2.0, 0.001, 0.2), ] - assert sampled_joint_path_frames([empty_position], duration=1.0, fps=10.0) == [] - assert ( - PreviewAnimator(set_calls.append, sleep=sleep_calls.append).animate( - [empty_position], duration=1.0, fps=10.0 - ) - is False - ) - assert set_calls == [] - assert sleep_calls == [] - - -def test_adapter_copies_joint_state_and_delegates_to_module() -> None: - copied = FakeJointState(["j1"], position=[1.0], velocity=[2.0], effort=[3.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(), None)}, - plan_to_pose=lambda pose, robot_name=None: (pose, robot_name), - plan_to_joints=lambda joints, robot_name=None: (joints, robot_name), - preview_path=lambda robot_name=None: robot_name, - execute=lambda robot_name=None: robot_name, - cancel=lambda: True, - clear_planned_path=lambda: True, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: copied, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: (robot_id, joint_state), - ) - module._world_monitor = world_monitor - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - current = adapter.get_current_joint_state("arm") - assert current is not copied - assert current.name is not copied.name - assert adapter.plan_to_pose("pose", "arm") == ("pose", "arm") - assert adapter.preview_path("arm") == "arm" - assert adapter.evaluate_joint_target(copied, "arm")["status"] == "FEASIBLE" +def test_init_and_home_presets_use_operator_init_and_config_home( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel([selected], states("arm")) + module.configs["arm"].home_joints = [0.9, 0.8] + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id].position == [-0.5, -1.0] -def test_adapter_evaluate_joint_target_uses_world_monitor_and_copies_input() -> None: - original = FakeJointState(["arm/j1", "j2"], position=[1.0, 2.0]) - seen = {} + gui._apply_preset("Home") + assert gui.state.group_joint_targets[selected.id].position == [0.9, 0.8] - def is_state_valid(robot_id, joint_state) -> bool: - seen["robot_id"] = robot_id - seen["joint_state"] = joint_state - return True - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: None, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=is_state_valid, - get_ee_pose=lambda robot_id, joint_state=None: (robot_id, joint_state), +def test_initial_pose_targets_are_group_id_keyed_for_same_robot_groups() -> None: + first = group("arm", "wrist", ("j1",), pose=True) + second = group("arm", "tool", ("j2",), pose=True) + module = Module([first, second], states("arm")) + monitor = Monitor(module) + monitor.poses[first.id] = Pose( + {"position": [1.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(), None)}, - _world_monitor=world_monitor, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - - result = adapter.evaluate_joint_target(original, "arm") - - assert result["success"] is True - assert result["status"] == "FEASIBLE" - assert seen["robot_id"] == "robot-1" - assert seen["joint_state"] is not original - assert seen["joint_state"].name == ["arm/j1", "j2"] - assert seen["joint_state"].position == [1.0, 2.0] - - -def test_obstacle_collision_marks_joint_target_infeasible() -> None: - obstacle = SimpleNamespace(name="blocking_box", blocked_joint_min=0.5) - - def is_state_valid(robot_id, joint_state) -> bool: - return bool(joint_state.position[0] < obstacle.blocked_joint_min) - - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: FakeJointState(["j1"], position=[0.0]), - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=is_state_valid, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), + monitor.poses[second.id] = Pose( + {"position": [2.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(joint_names=["j1"]), None)}, - _world_monitor=world_monitor, + server = Server() + scene_info = PlanningSceneInfo( + robots={"id-arm": module.configs["arm"]}, planning_groups=tuple(module.groups) ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - - free = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[0.25]), "arm") - colliding = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[0.75]), "arm") - - assert free["success"] is True - assert free["status"] == "FEASIBLE" - assert free["collision_free"] is True - assert colliding["success"] is True - assert colliding["status"] == "COLLISION" - assert colliding["collision_free"] is False - - -def test_scene_registers_goal_robot_coloring_and_updates_visibility() -> None: - server = FakeServer() - scene = ViserManipulationScene( + current = {"id-arm": JointState(module.states["arm"])} + gui = ViserPanelGui( server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, - ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], + scene_info, + Operator(module, monitor), + current, + ViserVisualizationConfig(panel_enabled=True), ) + try: + gui.start() + gui._toggle_group_selected(second.id) - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] + assert list(gui.state.pose_targets[first.id].position) == [1.0, 0.0, 0.0] + assert list(gui.state.pose_targets[second.id].position) == [2.0, 0.0, 0.0] + finally: + gui.close() - assert all(mesh.color == (255, 122, 0) for mesh in target._meshes) - assert all(mesh.opacity == 0.7 for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - assert all(mesh.opacity == 0.55 for mesh in preview._meshes) - scene.show_preview("robot1") - assert all(mesh.visible is True for mesh in preview._meshes) - scene.hide_preview("robot1") - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) +def test_panel_action_controls_are_present_in_source_order( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + _gui, _module, server = panel([selected], states("arm")) + + assert [button.label for button in server.gui.buttons[1:]] == [ + "Plan", + "Preview", + "Execute", + "Cancel", + "Clear plan", + ] + assert [folder.label for folder in server.gui.folders] == [ + "Manipulation Panel", + "Joint Control", + ] -def test_scene_transform_controls_update_pose_callback_and_visual_state() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( - server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, - ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], - ) - scene.register_robot("robot1", config) - updates = [] - - control = scene.ensure_target_controls("robot1", updates.append) - assert control is not None - assert server.transform_controls[0].path == "/targets/robot1/ee_control" - assert control.update_callback is not None - moved = SimpleNamespace(position=(1.0, 2.0, 3.0), wxyz=(1.0, 0.0, 0.0, 0.0)) - control.update_callback(SimpleNamespace(target=moved)) - assert updates == [moved] - - pose = Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}) - scene.set_target_pose("robot1", pose) - assert control.position == (0.1, 0.2, 0.3) - assert control.wxyz == (1.0, 0.0, 0.0, 0.0) - - scene.set_target_visual_state("robot1", feasible=False) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - assert control.color == (255, 40, 40) - assert all(mesh.color == (255, 30, 30) for mesh in target._meshes) - assert all(mesh.opacity == 0.75 for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - - -def test_scene_target_controls_update_target_ghost_pose_and_feasibility() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( - server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, +def test_target_callbacks_require_current_target_identity( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + first, second = ( + group("arm", "manipulator", ("j1",), pose=True), + group("arm", "gripper", ("j2",)), ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], + gui, _module, _server = panel([first, second], states("arm")) + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=(second.id,), ) - scene.register_robot("robot1", config) - scene.ensure_target_controls("robot1", lambda target: None) - pose = Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}) - assert scene.set_target_joints("robot1", ["joint1", "joint2"], [0.7, 0.9]) is True - assert scene.set_target_pose("robot1", pose) is None - assert scene.set_target_visual_state("robot1", feasible=False) is None + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) - target = scene._urdfs["robot1:target"] - handle = scene._handles["robot1:ee_control"] - assert target.cfg == [0.7, 0.9] - assert handle.position == (0.1, 0.2, 0.3) - assert handle.color == (255, 40, 40) + assert gui.state.target_status == TargetStatus.CHECKING -def test_gui_initializes_pose_selector_to_current_ee_pose( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1"], position=[0.25]) - current_pose = SimpleNamespace( - position=SimpleNamespace(x=0.1, y=0.2, z=0.3), - orientation=SimpleNamespace(w=0.9, x=0.1, y=0.2, z=0.3), - ) - config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - get_ee_pose=lambda robot_id, joint_state=None: current_pose, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - scene = ViserManipulationScene( - FakeTransformServer(), lambda *args, **kwargs: FakeViserUrdfWithMeshes(), preview_fps=10.0 - ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - control = scene._handles["robot-1:ee_control"] - assert control.position == (0.1, 0.2, 0.3) - assert control.wxyz == (0.9, 0.1, 0.2, 0.3) - assert gui.state.cartesian_target is current_pose +def test_scene_target_ghost_tracks_current_only_until_explicit_target() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene.update_current_robot("id-arm", JointState({"name": ["j1", "j2"], "position": [0.1, 0.2]})) + assert scene._urdfs["id-arm:target"].cfg == [0.1, 0.2] + scene.set_target_joints("id-arm", ["j1", "j2"], [0.8, 0.9]) + scene.update_current_robot("id-arm", JointState({"name": ["j1", "j2"], "position": [0.2, 0.3]})) -def test_gui_preset_dropdown_and_controls_include_init_home_current_and_callbacks( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["arm/j1", "arm/j2"], position=[0.25, 0.5]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[1.0, 2.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - _init_joints={"arm": FakeJointState(["j1", "j2"], position=[-1.0, -2.0])}, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - gui = make_panel(FakeGuiServer(), adapter) - assert gui._handles["preset"].options == ["Select preset...", "Init", "Current", "Home"] - assert list(gui._joint_sliders) == ["j1", "j2"] - gui._apply_preset("Home") - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [1.0, 2.0] - gui._apply_preset("Current") - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [0.25, 0.5] - gui._submit_execute() - assert "Cannot execute" in gui.state.error - - -def test_gui_rebuilding_joint_sliders_removes_stale_viser_handles( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[1.0, 2.0]) - module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - server = FakeGuiServer() - gui = make_panel(server, adapter) - stale_sliders = list(server.sliders) - assert [slider.value for slider in stale_sliders] == [0.0, 0.0] - - current.position = [-0.738, -0.2826151825863572] - gui._build_joint_sliders() - - assert all(slider.removed is True for slider in stale_sliders) - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [ - -0.738, - -0.2826151825863572, - ] + assert scene._urdfs["id-arm:target"].cfg == [0.8, 0.9] -def test_gui_parses_numpy_transform_control_arrays() -> None: - gui = ViserPanelGui(FakeGuiServer(), make_adapter_with_robot(), ViserVisualizationConfig()) +def test_scene_target_feasibility_colors_ghost_and_gizmo() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene.ensure_target_controls("id-arm", lambda _target: None) - pose = gui._pose_from_transform_target( - SimpleNamespace( - position=np.array([1.0, 2.0, 3.0]), - wxyz=np.array([0.5, 0.1, 0.2, 0.3]), - ) - ) + scene.set_target_visual_state("id-arm", False) - assert pose is not None - assert list(pose.position) == [1.0, 2.0, 3.0] - assert list(pose.orientation) == [0.1, 0.2, 0.3, 0.5] + assert scene._urdfs["id-arm:target"]._meshes[0].color == (255, 30, 30) + assert scene._handles["id-arm:ee_control"].color == (255, 40, 40) -def test_panel_execution_requires_fresh_plan_and_refresh_updates_robot_controls( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1"], position=[1.2]) - config = make_robot_config(joint_names=["j1"], home_joints=[0.5]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - execute=lambda robot_name=None: False, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, +def test_panel_feasibility_colors_group_controls_and_deduplicated_robot_ghosts() -> None: + arm_primary, arm_secondary, other = ( + group("arm", "primary", ("j1",), pose=True), + group("arm", "secondary", ("j2",), pose=True), + group("other", "manipulator", ("j1",), pose=True), ) - adapter = InProcessViserAdapter( - world_monitor=world_monitor, - manipulation_module=module, + module = Module([arm_primary, arm_secondary, other], states("arm", "other")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + scene.register_robot("id-arm", module.configs["arm"]) + scene.register_robot("id-other", module.configs["other"]) + gui = scene_gui(module, server, scene) + gui.start() + gui._worker.stop() + gui._operation_worker.stop() + gui._toggle_group_selected(arm_secondary.id) + gui._toggle_group_selected(other.id) + + control_calls: list[str] = [] + robot_calls: list[str] = [] + original_control = scene.set_target_control_visual_state + original_robot = scene.set_target_robot_visual_state + scene.set_target_control_visual_state = lambda group_id, feasible: ( + control_calls.append(group_id), + original_control(group_id, feasible), + ) # type: ignore[method-assign] + scene.set_target_robot_visual_state = lambda robot_id, feasible: ( + robot_calls.append(robot_id), + original_robot(robot_id, feasible), + ) # type: ignore[method-assign] + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, ) - gui = make_panel(FakeGuiServer(), adapter) - gui.refresh() - assert gui.state.selected_robot == "arm" - assert list(gui._joint_sliders) == ["j1"] - gui._apply_preset("Home") - assert gui._joint_sliders["j1"].value == 0.5 - gui._submit_execute() - assert "Cannot execute" in gui.state.error + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) - -def test_gui_moves_joint_target_immediately_and_stores_evaluated_joint_solution( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - target_pose = SimpleNamespace(position=SimpleNamespace(x=0.2, y=0.3, z=0.4)) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: target_pose, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - target_updates = [] - target_pose_updates = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: target_updates.append(args) or True, - set_target_pose=lambda *args: target_pose_updates.append(args), - set_target_visual_state=lambda *args: None, + assert control_calls == [arm_primary.id, arm_secondary.id, other.id] * 2 + assert robot_calls == ["id-arm", "id-other"] * 2 + assert all( + scene._handles[f"{item.id}:ee_control"].color == TARGET_CONTROL_FEASIBLE_COLOR + for item in (arm_primary, arm_secondary, other) ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - requests = [] - gui._worker.stop() - gui._worker = SimpleNamespace( - submit=lambda request: requests.append(request), stop=lambda: None + assert scene._urdfs["id-arm:target"]._meshes[0].color == GOAL_ROBOT_FEASIBLE_COLOR + assert scene._urdfs["id-other:target"]._meshes[0].color == GOAL_ROBOT_FEASIBLE_COLOR + + control_calls.clear() + robot_calls.clear() + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, ) - gui._joint_sliders["j1"].value = 0.25 - gui._joint_sliders["j2"].value = 0.75 - gui._submit_joint_target_evaluation() - assert target_updates[-1] == ("robot-1", ["j1", "j2"], [0.25, 0.75]) - assert target_pose_updates[-1] == ("robot-1", target_pose) - assert requests[-1].source == "joints" - - stale_request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - fresh_request = TargetEvaluationRequest(sequence_id=2, source="joints", robot_name="arm") - gui.state.latest_sequence_id = 2 gui._apply_target_evaluation_result( - stale_request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [9.0, 9.0]), - }, + request, TargetEvaluationResult(True, "COLLISION", "", False) ) - assert gui.state.joint_target == [0.25, 0.75] - gui._apply_target_evaluation_result( - fresh_request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [1.0, 2.0]), - }, + assert control_calls == [arm_primary.id, arm_secondary.id, other.id] * 2 + assert robot_calls == ["id-arm", "id-other"] * 2 + assert all( + scene._handles[f"{item.id}:ee_control"].color == TARGET_CONTROL_INFEASIBLE_COLOR + for item in (arm_primary, arm_secondary, other) ) - assert gui.state.target_status == TargetStatus.FEASIBLE - assert gui.state.feasibility.status == FeasibilityStatus.FEASIBLE - assert gui.state.joint_target == [1.0, 2.0] - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [0.25, 0.75] - assert target_updates[-1] == ("robot-1", ["j1", "j2"], [0.25, 0.75]) + assert scene._urdfs["id-arm:target"]._meshes[0].color == GOAL_ROBOT_INFEASIBLE_COLOR + assert scene._urdfs["id-other:target"]._meshes[0].color == GOAL_ROBOT_INFEASIBLE_COLOR + gui.close() -def test_gui_cartesian_ik_result_does_not_rewrite_active_gizmo( - make_panel: Callable[..., ViserPanelGui], +def test_scene_shared_clock_uses_stored_unequal_robot_frames( + monkeypatch: pytest.MonkeyPatch, ) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("left", config) + scene.register_robot("right", config) + updates: list[tuple[str, list[float]]] = [] + monkeypatch.setattr( + scene, + "_set_preview_ghost_joints", + lambda robot, _names, values: updates.append((robot, list(values))), ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - target_joint_updates = [] - target_pose_updates = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: target_joint_updates.append(args) or True, - set_target_pose=lambda *args: target_pose_updates.append(args), - set_target_visual_state=lambda *args: None, + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", lambda _delay: None ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - gui.state.cartesian_target = Pose( - {"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]} + + assert scene.animate_preview( + GroupPreviewAnimation( + ( + PreviewTrack("left", ("j1",), (PreviewFrame(0.0, (0.0,)),)), + PreviewTrack( + "right", + ("j1",), + ( + PreviewFrame(0.0, (10.0,)), + PreviewFrame(1.0, (11.0,)), + ), + ), + ) + ), + 1.0, ) - request = TargetEvaluationRequest(sequence_id=1, source="cartesian", robot_name="arm") - gui.state.latest_sequence_id = 1 + assert updates == [ + ("left", [0.0]), + ("right", [10.0]), + ("left", [0.0]), + ("right", [11.0]), + ] - gui._apply_target_evaluation_result( - request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [1.0, 2.0]), - }, + +def test_animation_frame_helpers_scale_stored_timestamps() -> None: + frames = ( + PreviewFrame(0.0, (0.0,)), + PreviewFrame(0.25, (1.0,)), + PreviewFrame(1.0, (2.0,)), ) - assert gui.state.target_status == TargetStatus.FEASIBLE - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [1.0, 2.0] - assert target_joint_updates[-1] == ("robot-1", ["j1", "j2"], [1.0, 2.0]) - assert target_pose_updates == [] + assert scaled_frame_delays(frames, 2.0) == (0.5, 1.5) -def test_gui_collision_evaluation_marks_target_infeasible_and_colors_scene( - make_panel: Callable[..., ViserPanelGui], +def test_panel_disables_plan_preview_and_execute_until_a_feasible_target( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["j1"], position=[0.0]) - config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule(_robots={"arm": ("robot-1", config, None)}) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: False, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), - ) - module._world_monitor = world_monitor - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - visual_states = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: True, - set_target_pose=lambda *args: None, - set_target_visual_state=lambda *args: visual_states.append(args), - ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - gui.state.latest_sequence_id = 1 - result = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[1.0]), "arm") + selected = group("arm", "manipulator", ("j1",), pose=True) + _gui, _module, server = panel([selected], states("arm")) - gui._apply_target_evaluation_result(request, result) + assert [button.disabled for button in server.gui.buttons[1:4]] == [True, True, True] - assert result["status"] == "COLLISION" - assert gui.state.target_status == TargetStatus.INFEASIBLE - assert gui.state.feasibility.status == FeasibilityStatus.COLLISION - assert gui.state.error == "Target is in collision" - assert visual_states[-1] == ("robot-1", False) - -def test_gui_safe_execute_requires_fresh_matching_plan_and_clear_resets_path( - make_panel: Callable[..., ViserPanelGui], monkeypatch: pytest.MonkeyPatch +def test_panel_status_reports_target_and_plan_defaults( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["j1"], position=[1.0]) - executed = [] - cleared = [] - module = FakeManipulationModule( - _robots={ - "arm": ("robot-1", make_robot_config(joint_names=["j1"], home_joints=[1.0]), None) - }, - _state=NamedState(name="IDLE"), - execute=lambda robot_name=None: executed.append(robot_name) or True, - clear_planned_path=lambda: cleared.append(True) or True, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - gui = make_panel( - FakeGuiServer(), - adapter, - ViserVisualizationConfig(panel_enabled=True, current_match_tolerance=0.05), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), - ) - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.plan_state = PanelPlanState( - status=PlanStatus.FRESH, - robot="arm", - start_joints_snapshot=[1.2], - ) - gui._submit_execute() - assert executed == [] - assert "Cannot execute" in gui.state.error + selected = group("arm", "manipulator", ("j1",), pose=True) + gui, _module, server = panel([selected], states("arm")) - gui.state.action_status = ActionStatus.IDLE - gui.state.error = "" - gui.state.plan_state.start_joints_snapshot = [1.0] - gui._submit_execute() - assert executed == ["arm"] + assert server.gui.markdown[0].value == "### Status\n**State:** Ready" + assert gui.state.error == "" - gui._submit_clear() - assert cleared == [True] - assert gui.state.plan_state.status == PlanStatus.NONE +def test_scene_reference_grid_has_expected_defaults_and_toggle() -> None: + server = Server() + grids: list[Handle] = [] + server.scene.add_grid = lambda *_args, **_kwargs: grids.append(Handle()) or grids[-1] + scene = ViserManipulationScene(server, Urdf) -def test_gui_plan_target_failure_recovers_action_state( - make_panel: Callable[..., ViserPanelGui], - monkeypatch: pytest.MonkeyPatch, -) -> None: - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), - ) - gui.state.selected_robot = "missing" - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "IDLE" - - gui._submit_plan() + assert scene.has_reference_grid() is True + scene.set_reference_grid_visible(False) + assert grids[0].visible is False + scene.set_reference_grid_visible(True) + assert grids[0].visible is True - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.plan_state.status == PlanStatus.FAILED - assert gui.state.error == "No robot config" - assert gui.state.last_result == "plan_to_joints=False" +def test_scene_returns_false_for_missing_robot_target_updates() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) -def test_gui_resets_fault_before_replanning( - make_panel: Callable[..., ViserPanelGui], - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = [] - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), + assert scene.set_target_joints("missing", ["j1"], [0.1]) is False + assert ( + scene.animate_preview( + GroupPreviewAnimation( + (PreviewTrack("missing", ("j1",), (PreviewFrame(0.0, (0.0,)),)),) + ), + duration=0.0, + ) + is False ) - def reset() -> bool: - calls.append("reset") - return True - def plan_to_joints(_joints: JointState, _robot_name: str | None = None) -> bool: - calls.append("plan") - return True +def test_scene_cancel_generation_hides_preview_and_rejects_old_animation() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene._preview_visible["id-arm"] = True + scene._set_preview_visibility("id-arm", True) - monkeypatch.setattr(adapter, "reset", reset) - monkeypatch.setattr(adapter, "plan_to_joints", plan_to_joints) - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "FAULT" - - gui._submit_plan() + scene.cancel_preview_animation() - assert calls == ["reset", "plan"] - assert gui.state.plan_state.status == PlanStatus.FRESH - assert gui.state.last_result == "plan_to_joints=True" + assert scene._preview_visible == {"id-arm": False} + assert scene._animation_generation == 1 -def test_operation_worker_coalesces_pending_requests() -> None: - errors = [] - calls = [] - worker = OperationWorker(errors.append) - worker.submit(lambda: calls.append("old")) - worker.submit(lambda: calls.append("new")) +def test_scene_base_pose_requires_urdf_root_to_match(monkeypatch: pytest.MonkeyPatch) -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.parse_model", + lambda _path: SimpleNamespace(root_link="world"), + ) - operation = worker._requests.get_nowait() - operation.operation() + with pytest.raises(ValueError, match="base_link 'base'.*URDF root 'world'"): + scene._assert_base_link_is_urdf_root(SimpleNamespace(base_link="base"), "robot.urdf") - assert calls == ["new"] - assert errors == [] +def test_scene_detects_non_identity_base_pose() -> None: + identity = SimpleNamespace( + base_pose=Pose({"position": [0.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]}) + ) + translated = SimpleNamespace( + base_pose=Pose({"position": [1.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]}) + ) -def test_operation_worker_stop_can_wait_for_in_flight_operation() -> None: - errors = [] - worker = OperationWorker(errors.append) - started = threading.Event() - release = threading.Event() - finished = threading.Event() - stopped = threading.Event() + assert ViserManipulationScene._has_non_identity_base_pose(identity) is False + assert ViserManipulationScene._has_non_identity_base_pose(translated) is True - def operation() -> None: - started.set() - release.wait(timeout=1.0) - finished.set() - worker.start() - worker.submit(operation) - assert started.wait(timeout=1.0) +@pytest.mark.parametrize("interruption", ["cancel", "replacement", "close"]) +def test_scene_inflight_preview_never_updates_after_generation_replacement( + monkeypatch: pytest.MonkeyPatch, interruption: str +) -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + first_tick, release = threading.Event(), threading.Event() + updates: list[float] = [] + original = scene._set_preview_ghost_joints + + def record(robot_id: str, names: Sequence[str], values: Sequence[float]) -> None: + updates.append(float(values[0])) + original(robot_id, names, values) + + scene._set_preview_ghost_joints = record # type: ignore[method-assign] + sleep_calls = 0 + + def block_after_first_tick(_delay: float) -> None: + nonlocal sleep_calls + sleep_calls += 1 + if sleep_calls == 1: + first_tick.set() + assert release.wait(timeout=2.0) - stopper = threading.Thread( - target=lambda: (worker.stop(timeout=None), stopped.set()), - name="StopViserOperationTest", + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", block_after_first_tick + ) + old = GroupPreviewAnimation( + ( + PreviewTrack( + "id-arm", + ("j1",), + ( + PreviewFrame(0.0, (0.0,)), + PreviewFrame(1.0, (1.0,)), + ), + ), + ) ) - stopper.start() - assert not stopped.wait(timeout=0.05) + worker = threading.Thread(target=lambda: scene.animate_preview(old, 1.0)) + worker.start() + assert first_tick.wait(timeout=2.0) + if interruption == "cancel": + visualizer = ViserManipulationVisualizer( + config=ViserVisualizationConfig(panel_enabled=False), + ) + visualizer._scene = scene + visualizer.cancel_preview_animation() + stable_updates = list(updates) + elif interruption == "replacement": + updates.clear() + assert scene.animate_preview( + GroupPreviewAnimation( + (PreviewTrack("id-arm", ("j1",), (PreviewFrame(0.0, (10.0,)),)),) + ), + 0.0, + ) + stable_updates = list(updates) + else: + scene.close() + stable_updates = list(updates) release.set() - assert stopped.wait(timeout=1.0) - stopper.join(timeout=1.0) - - assert finished.is_set() - assert worker._thread is None - assert errors == [] - + worker.join(timeout=2.0) -def test_target_evaluation_worker_coalesces_pending_requests() -> None: - worker = TargetEvaluationWorker(lambda request: {}, lambda request, result: None) - old_request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - new_request = TargetEvaluationRequest(sequence_id=2, source="joints", robot_name="arm") + assert not worker.is_alive() + assert updates == stable_updates - worker.submit(old_request) - worker.submit(new_request) - - assert worker._requests.get_nowait() is new_request +def test_transform_control_callback_preserves_pose_through_gui_and_backend( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + module = Module([selected], states("arm")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + gui = scene_gui(module, server, scene) + submitted: list[TargetEvaluationRequest] = [] + gui._worker.submit = submitted.append # type: ignore[method-assign] + gui.start() + control = scene._handles[f"{selected.id}:ee_control"] + control.position = (1.0, 2.0, 3.0) + control.wxyz = (0.4, 0.1, 0.2, 0.3) + assert control.callback is not None + control.callback(SimpleNamespace(target=control)) + + request = submitted[-1] + assert list(gui.state.pose_targets[selected.id].position) == [1.0, 2.0, 3.0] + assert list(gui.state.pose_targets[selected.id].orientation) == [0.1, 0.2, 0.3, 0.4] + assert control.position == (1.0, 2.0, 3.0) + assert control.wxyz == (0.4, 0.1, 0.2, 0.3) + assert request.pose_targets[selected.id] == gui.state.pose_targets[selected.id] + gui.close() -def test_operation_worker_reports_timeout() -> None: - errors = [] - release = threading.Event() - finished = threading.Event() - worker = OperationWorker(errors.append, timeout_seconds=0.01) - def operation() -> None: - release.wait(timeout=1.0) - finished.set() +def test_joint_evaluation_updates_active_gizmo_from_computed_group_pose() -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + module = Module([selected], states("arm")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + gui = scene_gui(module, server, scene) + gui.start() + control = scene._handles[f"{selected.id}:ee_control"] + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, + ) + computed_pose = Pose({"position": [0.7, 0.8, 0.9], "orientation": [0.1, 0.2, 0.3, 0.4]}) - worker.submit(operation, timeout_seconds=0.01) - worker._run_operation(worker._requests.get_nowait()) - release.set() + gui._apply_target_evaluation_result( + request, + TargetEvaluationResult( + True, "FEASIBLE", "", True, group_poses={selected.id: computed_pose} + ), + ) - assert errors == ["Operation timed out after 0.0s"] - assert finished.wait(timeout=1.0) + assert control.position == (0.7, 0.8, 0.9) + assert control.wxyz == (0.4, 0.1, 0.2, 0.3) + gui.close() diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index 63891bad25..58fc0d9da1 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -23,17 +23,22 @@ from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) from dimos.manipulation.visualization.viser import ( runtime as runtime_module, visualizer as visualizer_module, ) -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.runtime import ViserRuntime +from dimos.manipulation.visualization.viser.scene import ViserManipulationScene from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class FakeDependency: @@ -79,8 +84,6 @@ def fail_runtime(_config: ViserVisualizationConfig) -> FakeServer: monkeypatch.setattr(visualizer_module, "ViserRuntime", fail_runtime) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=False), ) @@ -111,8 +114,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: calls.append(("create", "scene")) @@ -126,7 +127,9 @@ class FakeGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -146,8 +149,6 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FakeGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) scene = PlanningSceneInfo( @@ -157,7 +158,7 @@ def close(self) -> None: } ) - visualizer.initialize_scene(scene) + visualizer.initialize(VisualizationSession(scene, operator=FakeDependency())) assert calls == [ ("start", "runtime"), @@ -192,8 +193,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: pass @@ -204,7 +203,9 @@ class FakeGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -221,13 +222,13 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FakeGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) with pytest.raises(RuntimeError, match="gui failed"): - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) assert closed == ["gui", "scene", "runtime"] assert visualizer.get_visualization_url() is None @@ -255,8 +256,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: raise RuntimeError("scene failed") @@ -264,13 +263,13 @@ def __init__( monkeypatch.setattr(visualizer_module, "ViserUrdf", FakeViserUrdf) monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FailingScene) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=False), ) with pytest.raises(RuntimeError, match="scene failed"): - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) assert closed == ["runtime"] assert visualizer.get_visualization_url() is None @@ -298,8 +297,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: pass @@ -310,7 +307,9 @@ class FailingGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -331,11 +330,11 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FailingGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) with pytest.raises(RuntimeError, match="gui close failed"): visualizer.close() @@ -397,8 +396,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: calls.append(("scene", "create")) @@ -406,48 +403,122 @@ def update_current_robot(self, robot_id: str, joint_state: JointState | None) -> assert joint_state == current calls.append(("update", robot_id)) - def show_preview(self, robot_id: str) -> None: - calls.append(("show", robot_id)) + def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: + calls.append(("register", robot_id)) - def hide_preview(self, robot_id: str) -> None: - calls.append(("hide", robot_id)) + def cancel_preview_animation(self) -> None: + calls.append(("cancel", "preview")) - def animate_path(self, robot_id: str, path: list[JointState], duration: float) -> None: - assert path == [current] + def animate_preview(self, preview: object, duration: float) -> None: assert duration == 1.5 - calls.append(("animate", robot_id)) + calls.append(("animate", "groups")) def close(self) -> None: calls.append(("scene", "close")) - world_monitor = SimpleNamespace(get_current_joint_state=lambda _robot_id: current) - manipulation_module = SimpleNamespace( - robot_items=lambda: [("arm", "robot-1", fake_robot_config("arm"))], - robot_id_for_name=lambda robot_name: "robot-1" if robot_name == "arm" else None, - ) monkeypatch.setattr(visualizer_module, "ViserRuntime", FakeRuntime) monkeypatch.setattr(visualizer_module, "ViserUrdf", FakeViserUrdf) monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) visualizer = ViserManipulationVisualizer( - world_monitor=world_monitor, - manipulation_module=manipulation_module, config=ViserVisualizationConfig(panel_enabled=False), ) - visualizer.publish_visualization() - visualizer.show_preview("robot-1") - visualizer.hide_preview("robot-1") - visualizer.animate_path("robot-1", [current], duration=1.5) + assert hasattr(ViserManipulationVisualizer, "cancel_preview_animation") + visualizer.initialize( + VisualizationSession(PlanningSceneInfo({"robot-1": fake_robot_config("arm")})) + ) + visualizer.cancel_preview_animation() + visualizer.update_state(VisualizationStateFrame({"robot-1": current})) + visualizer.cancel_preview_animation() + visualizer.animate_trajectory(JointTrajectory(joint_names=["arm/joint1"]), duration=1.5) visualizer.close() - visualizer.publish_visualization() + visualizer.update_state(VisualizationStateFrame({"robot-1": current})) assert calls == [ ("runtime", "start"), ("scene", "create"), + ("register", "robot-1"), + ("cancel", "preview"), ("update", "robot-1"), - ("show", "robot-1"), - ("hide", "robot-1"), - ("animate", "robot-1"), + ("cancel", "preview"), + ("animate", "groups"), ("scene", "close"), ("runtime", "close"), ] + + +def test_scene_prepares_urdf_applies_base_pose_and_rejects_wrong_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + created: list[tuple[Path, str]] = [] + prepared: list[dict[str, object]] = [] + frames: list[dict[str, object]] = [] + fixed_world_root = tmp_path / "fixed-world.urdf" + fixed_world_root.write_text( + """ + + +""" + ) + non_fixed_world_root = tmp_path / "non-fixed-world.urdf" + non_fixed_world_root.write_text( + """ + + +""" + ) + + class Handle: + def remove(self) -> None: + return None + + class SceneApi: + def add_frame(self, name: str, **kwargs: object) -> Handle: + frames.append({"name": name, **kwargs}) + return Handle() + + class Server: + scene = SceneApi() + + class Urdf: + def __init__( + self, _server: Server, path: Path, *, root_node_name: str, **_kwargs: object + ) -> None: + created.append((path, root_node_name)) + self._meshes: list[object] = [] + + config = fake_robot_config("arm") + config.base_pose.position.x = 1.0 + + def prepare(path: Path, **kwargs: object) -> Path: + prepared.append(kwargs) + return fixed_world_root if path.name == "arm.urdf" else non_fixed_world_root + + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.prepare_urdf_for_drake", + prepare, + ) + + def parse_prepared_model(path: Path) -> SimpleNamespace: + content = path.read_text() + return SimpleNamespace(root_link="world" if 'name="world"' in content else "base_link") + + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.parse_model", parse_prepared_model + ) + scene = ViserManipulationScene(Server(), Urdf) + + scene.register_robot("robot-1", config) + + assert [root for _, root in created] == [ + "/robots/robot-1/current/base_pose/urdf", + "/targets/robot-1/target/base_pose/urdf", + "/previews/robot-1/ghost/base_pose/urdf", + ] + assert prepared == [{"package_paths": {}, "xacro_args": {}, "convert_meshes": False}] * 3 + assert all('name="world"' not in path.read_text() for path, _ in created) + assert all(frame["position"] == (1.0, 0.0, 0.0) for frame in frames) + wrong_root_config = fake_robot_config("wrong") + with pytest.raises(ValueError, match="prepared URDF root 'world'"): + scene.prepared_urdf_path(wrong_root_config) diff --git a/dimos/manipulation/visualization/viser/theme.py b/dimos/manipulation/visualization/viser/theme.py index 339be04492..d38767d1da 100644 --- a/dimos/manipulation/visualization/viser/theme.py +++ b/dimos/manipulation/visualization/viser/theme.py @@ -36,7 +36,7 @@ DIMOS_THEME_TITLE = "DimOS Manipulation" DIMOS_THEME_URL = "https://github.com/dimensionalOS/dimos" -DIMOS_BRAND_COLOR = (22, 130, 163) +DIMOS_BRAND_COLOR = (0, 153, 255) DIMOS_LOGO_PATH = Path(__file__).with_name("assets") / "dimensional-logo.svg" logger = setup_logger() @@ -86,8 +86,8 @@ def _configure_theme(server: ViserServer, titlebar_content: TitlebarConfig | Non try: server.gui.configure_theme( titlebar_content=titlebar_content, - control_layout="collapsible", - control_width="medium", + control_layout="fixed", + control_width="large", dark_mode=True, show_logo=False, show_share_button=False, diff --git a/dimos/manipulation/visualization/viser/visualizer.py b/dimos/manipulation/visualization/viser/visualizer.py index d12c98e53b..a9e6ee4c85 100644 --- a/dimos/manipulation/visualization/viser/visualizer.py +++ b/dimos/manipulation/visualization/viser/visualizer.py @@ -14,10 +14,15 @@ from __future__ import annotations +from collections.abc import Sequence from contextlib import suppress from typing import TYPE_CHECKING -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from dimos.manipulation.visualization.viser.animation import ( + GroupPreviewAnimation, + PreviewFrame, + PreviewTrack, +) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.gui import ViserPanelGui from dimos.manipulation.visualization.viser.runtime import ( @@ -27,6 +32,8 @@ ) from dimos.manipulation.visualization.viser.scene import ViserManipulationScene from dimos.manipulation.visualization.viser.theme import apply_dimos_theme +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger try: @@ -41,12 +48,11 @@ raise ModuleNotFoundError(VISER_URDF_INSTALL_HINT) from e if TYPE_CHECKING: - from dimos.manipulation.manipulation_module import ManipulationModule - from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor + from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( - JointPath, PlanningSceneInfo, - WorldRobotID, + VisualizationSession, + VisualizationStateFrame, ) logger = setup_logger() @@ -58,18 +64,19 @@ class ViserManipulationVisualizer: def __init__( self, *, - world_monitor: WorldMonitor, - manipulation_module: ManipulationModule, config: ViserVisualizationConfig | None = None, ) -> None: - self._world_monitor = world_monitor - self._manipulation_module = manipulation_module self.config = config or ViserVisualizationConfig() self._runtime: ViserRuntime | None = None self._server: ViserServer | None = None - self._adapter: InProcessViserAdapter | None = None self._scene: ViserManipulationScene | None = None self._gui: ViserPanelGui | None = None + self._session_scene: PlanningSceneInfo | None = None + self._operator: object | None = None + self._current_states: dict[str, JointState] = {} + self._robot_names_by_id: dict[str, str] = {} + self._robot_ids_by_name: dict[str, str] = {} + self._configs_by_name: dict[str, RobotModelConfig] = {} self._closed = False def _ensure_started(self) -> None: @@ -81,18 +88,19 @@ def _ensure_started(self) -> None: try: server = runtime.start() apply_dimos_theme(server) - adapter = InProcessViserAdapter( - world_monitor=self._world_monitor, - manipulation_module=self._manipulation_module, - ) - scene = ViserManipulationScene( - server, - ViserUrdf, - preview_fps=self.config.preview_fps, - ) + scene = ViserManipulationScene(server, ViserUrdf) gui = ( - ViserPanelGui(server, adapter, self.config, scene) + ViserPanelGui( + server, + self._session_scene, + self._operator, + self._current_states, + self.config, + scene, + ) if self.config.panel_enabled + and self._session_scene is not None + and self._operator is not None else None ) if gui is not None: @@ -108,20 +116,31 @@ def _ensure_started(self) -> None: runtime.close() self._runtime = None self._server = None - self._adapter = None self._scene = None self._gui = None self._closed = True raise self._runtime = runtime self._server = server - self._adapter = adapter self._scene = scene self._gui = gui self._closed = False logger.info(f"Viser manipulation visualization: {self.get_visualization_url()}") - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: + """Initialize Viser robot visuals from a one-shot visualization session.""" + self._operator = session.operator + self._session_scene = session.scene + self._robot_names_by_id = { + str(robot_id): config.name for robot_id, config in session.scene.robots.items() + } + self._robot_ids_by_name = { + config.name: str(robot_id) for robot_id, config in session.scene.robots.items() + } + self._configs_by_name = {config.name: config for config in session.scene.robots.values()} + self._initialize_scene(session.scene) + + def _initialize_scene(self, scene: PlanningSceneInfo) -> None: """Initialize Viser robot visuals from planning-scene metadata.""" if self._closed: return @@ -140,45 +159,94 @@ def initialize_scene(self, scene: PlanningSceneInfo) -> None: def get_visualization_url(self) -> str | None: return None if self._runtime is None else self._runtime.url - def publish_visualization(self, ctx: None = None) -> None: - """Update current robot render state. ctx is accepted for protocol compatibility.""" + def update_state(self, frame: VisualizationStateFrame) -> None: + """Update current robot render state from a pushed state frame.""" if self._closed: return self._ensure_started() - if self._adapter is None or self._scene is None: + if self._scene is None: return - for _robot_name, robot_id, _config in self._adapter.robot_items(): - current = self._adapter.get_current_joint_state(_robot_name) - self._scene.update_current_robot(str(robot_id), current) + for robot_id, current in frame.joint_states.items(): + robot_id_string = str(robot_id) + self._current_states[robot_id_string] = JointState(current) + self._scene.update_current_robot(robot_id_string, current) if self._gui is not None: self._gui.refresh() - def show_preview(self, robot_id: WorldRobotID) -> None: - if not self._closed: - self._ensure_started() - if self._scene is None: - return - self._scene.show_preview(str(robot_id)) - - def hide_preview(self, robot_id: WorldRobotID) -> None: - if not self._closed: - self._ensure_started() - if self._scene is None: - return - self._scene.hide_preview(str(robot_id)) - - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: if self._closed: return self._ensure_started() if self._scene is None: return - self._scene.animate_path(str(robot_id), list(path), duration) + preview = self._raw_preview_animation(trajectory) + if preview is not None: + self._scene.animate_preview( + preview, duration if duration is not None else max(float(trajectory.duration), 0.0) + ) + + def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> None: + """Cancel preview playback without starting a renderer or waiting for it. + + The world monitor deliberately invokes this outside its visualization + lock, so a renderer sleeping between frames can observe the scene + generation change immediately. Do not call ``_ensure_started()``: + cancelling before Viser has started must remain a no-op. Likewise, + retain the scene reference while ``close()`` is unwinding so a + concurrent cancellation can still invalidate an in-flight frame. + """ + scene = self._scene + if scene is not None: + if robot_ids is None: + scene.cancel_preview_animation() + else: + scene.cancel_preview_animation(robot_ids) + + def _raw_preview_animation(self, trajectory: JointTrajectory) -> GroupPreviewAnimation | None: + robot_indices: dict[str, list[tuple[int, str]]] = {} + for index, global_name in enumerate(trajectory.joint_names): + if "/" not in str(global_name): + return None + robot_name, local_name = str(global_name).split("/", 1) + if robot_name not in self._robot_ids_by_name: + return None + robot_indices.setdefault(robot_name, []).append((index, local_name)) + tracks: list[PreviewTrack] = [] + for robot_name, indexed_names in robot_indices.items(): + robot_id = self._robot_ids_by_name[robot_name] + config = self._configs_by_name[robot_name] + current = self._current_states.get(robot_id) + baseline = self._baseline_values(config, current) + if baseline is None: + return None + frames: list[PreviewFrame] = [] + for point in trajectory.points: + selected = { + local_name: float(point.positions[index]) for index, local_name in indexed_names + } + positions: list[float] = [] + for local_name in config.joint_names: + value = selected.get(local_name, baseline.get(local_name)) + if value is None: + return None + positions.append(float(value)) + frames.append(PreviewFrame(float(point.time_from_start), tuple(positions))) + tracks.append(PreviewTrack(robot_id, tuple(config.joint_names), tuple(frames))) + return GroupPreviewAnimation(tuple(tracks)) if tracks else None + + @staticmethod + def _baseline_values( + config: RobotModelConfig, current: JointState | None + ) -> dict[str, float] | None: + if current is None or len(current.name) != len(current.position): + return None + values = { + str(name): float(value) + for name, value in zip(current.name, current.position, strict=True) + } + return values if all(name in values for name in config.joint_names) else None def close(self) -> None: if self._closed: @@ -204,7 +272,6 @@ def close(self) -> None: errors.append(e) self._runtime = None self._server = None - self._adapter = None self._scene = None self._gui = None if errors: diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 6b4dd8c68e..abd92725b3 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -60,6 +60,7 @@ "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic", "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic", "dual-xarm6-planner": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner", + "dual-xarm6-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator", "dual-xarm6-planner-coordinator-mock-meshcat": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator_mock_meshcat", "keyboard-teleop-a750": "dimos.robot.manipulators.a750.blueprints.teleop:keyboard_teleop_a750", "keyboard-teleop-openarm": "dimos.robot.manipulators.openarm.blueprints.teleop:keyboard_teleop_openarm", diff --git a/dimos/robot/manipulators/test_blueprints.py b/dimos/robot/manipulators/test_blueprints.py index c8bd7ddaa9..0720362e4f 100644 --- a/dimos/robot/manipulators/test_blueprints.py +++ b/dimos/robot/manipulators/test_blueprints.py @@ -21,9 +21,12 @@ MeshcatVisualizationConfig, NoManipulationVisualizationConfig, ) +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.robot.get_all_blueprints import get_blueprint_by_name from dimos.robot.manipulators.common.blueprints import planner from dimos.robot.manipulators.xarm.blueprints.basic import ( dual_xarm6_planner, + dual_xarm6_planner_coordinator, dual_xarm6_planner_coordinator_mock_meshcat, xarm6_planner_only, xarm7_planner_coordinator, @@ -65,15 +68,23 @@ def test_xarm_planner_blueprints_default_to_no_visualization() -> None: assert isinstance(config.visualization, NoManipulationVisualizationConfig) -def test_dual_xarm6_mock_meshcat_blueprint_wires_planner_and_coordinator() -> None: - config = _manipulation_config(dual_xarm6_planner_coordinator_mock_meshcat) +def test_dual_xarm6_planner_coordinator_blueprints_preserve_visualization_backends() -> None: + assert get_blueprint_by_name("dual-xarm6-planner-coordinator") is dual_xarm6_planner_coordinator + assert ( + get_blueprint_by_name("dual-xarm6-planner-coordinator-mock-meshcat") + is dual_xarm6_planner_coordinator_mock_meshcat + ) + + config = _manipulation_config(dual_xarm6_planner_coordinator) + legacy_config = _manipulation_config(dual_xarm6_planner_coordinator_mock_meshcat) coordinator_kwargs = next( atom.kwargs - for atom in dual_xarm6_planner_coordinator_mock_meshcat.blueprints + for atom in dual_xarm6_planner_coordinator.blueprints if atom.module is ControlCoordinator ) - assert isinstance(config.visualization, MeshcatVisualizationConfig) + assert isinstance(config.visualization, ViserVisualizationConfig) + assert isinstance(legacy_config.visualization, MeshcatVisualizationConfig) assert [robot.name for robot in config.robots] == ["left_arm", "right_arm"] assert [hardware.hardware_id for hardware in coordinator_kwargs["hardware"]] == [ "left_arm", diff --git a/dimos/robot/manipulators/xarm/blueprints/basic.py b/dimos/robot/manipulators/xarm/blueprints/basic.py index 68e6bb467e..c3f623356d 100644 --- a/dimos/robot/manipulators/xarm/blueprints/basic.py +++ b/dimos/robot/manipulators/xarm/blueprints/basic.py @@ -47,6 +47,23 @@ _mock_left_xarm6_hw = make_xarm_hardware("left_arm", 6) _mock_right_xarm6_hw = make_xarm_hardware("right_arm", 6) +dual_xarm6_planner_coordinator = autoconnect( + planner( + robots=[ + make_xarm6_model_config(name="left_arm", y_offset=0.5), + make_xarm6_model_config(name="right_arm", y_offset=-0.5), + ], + visualization={"backend": "viser"}, + ), + coordinator( + hardware=[_mock_left_xarm6_hw, _mock_right_xarm6_hw], + tasks=[ + trajectory_task(_mock_left_xarm6_hw), + trajectory_task(_mock_right_xarm6_hw), + ], + ), +) + dual_xarm6_planner_coordinator_mock_meshcat = autoconnect( planner( robots=[ diff --git a/dimos/robot/test_all_blueprints.py b/dimos/robot/test_all_blueprints.py index ba907315fe..c31d3c8663 100644 --- a/dimos/robot/test_all_blueprints.py +++ b/dimos/robot/test_all_blueprints.py @@ -50,6 +50,7 @@ "coordinator-xarm6", "coordinator-xarm7", "dual-xarm6-planner", + "dual-xarm6-planner-coordinator", "dual-xarm6-planner-coordinator-mock-meshcat", "learning-collect-quest-piper", "learning-collect-quest-xarm7", diff --git a/docs/capabilities/manipulation/planning_groups.md b/docs/capabilities/manipulation/planning_groups.md index 043899fd50..f2dac09afb 100644 --- a/docs/capabilities/manipulation/planning_groups.md +++ b/docs/capabilities/manipulation/planning_groups.md @@ -13,7 +13,7 @@ torso, without confusing that group with the robot's hardware identity. | Planning group ID | Stable API ID in the form `{robot_name}/{group_name}`. | | Local joint name | Joint name inside a robot model, such as `joint1`. | | Global joint name | Boundary-level joint name in the form `{robot_name}/{local_joint_name}`. | -| Generated plan | Planning artifact containing selected group IDs and one synchronized global-joint path. | +| Generated plan | Planning artifact containing selected group IDs, geometric waypoints, and one synchronized global-joint trajectory. | | Auxiliary group | A selected group that contributes free DOFs to a pose plan without receiving its own pose target. | Local URDF/SRDF joint names stay inside robot-scoped configuration, model @@ -129,15 +129,20 @@ partial joint sets. A `GeneratedPlan` stores: - selected planning group IDs; -- a single synchronized path of `JointState` waypoints keyed by global joint - names; +- a geometric path of `JointState` waypoints keyed by global joint names; +- one materialized synchronized `JointTrajectory` over the same selected global + joint names; - status, timing, path length, iteration count, and message metadata. -Preview and execution project this path lazily. Preview sends projected joint -paths to the world monitor. Execution splits the path by affected trajectory -task, orders each trajectory by the robot's configured local joint order, writes -global joint names at the coordinator boundary, and invokes each trajectory -controller. Controllers remain planning-group agnostic. +Preview and execution consume the stored trajectory; they do not lazily +parameterize the geometric path. Preview forwards the raw globally named +trajectory through the visualization boundary, where renderers project it to +their robot-local visuals while preserving stored timestamps. Execution splits +the stored trajectory by affected trajectory task, translates selected joint +names at the coordinator boundary, and invokes each trajectory controller +without filling or commanding omitted joints. Controllers remain planning-group +agnostic, and trajectory tasks still claim their full configured joint set while +executing only the active planned subset. Multi-task dispatch is not atomic: if one trajectory task accepts and a later task rejects, DimOS reports the rejection but does not roll back the accepted diff --git a/docs/capabilities/manipulation/readme.md b/docs/capabilities/manipulation/readme.md index 2d0aba6afb..82f071d1a1 100644 --- a/docs/capabilities/manipulation/readme.md +++ b/docs/capabilities/manipulation/readme.md @@ -148,8 +148,9 @@ Backend choices: - `meshcat`: embedded Drake/Meshcat visualizer. The planning world must be created with embedded visualization enabled, so this is selected through the visualization config. -- `viser`: in-process Viser visualizer. It renders current robot state, target controls, - transient preview ghosts, planned path previews, and optional panel controls. +- `viser`: in-process Viser visualizer. It renders pushed current robot state, + target controls, transient preview ghosts, synchronized trajectory previews, + and optional panel controls. - `none`: no manipulation planning visualization. CLI example: @@ -184,20 +185,27 @@ Viser support is included in the `manipulation` extra: uv sync --extra manipulation --inexact ``` -The Viser panel uses existing manipulation planning, preview, execute, cancel, and clear-plan -RPC methods through a small in-process adapter. GUI callbacks enqueue operations instead of -touching `WorldSpec`, IK, planner objects, or live Drake contexts directly. Rendering copies -mutable joint state/path containers at the read boundary, then updates the Viser scene after -manipulation/world accessors have returned. - -External manipulation visualizers are initialized from a backend-neutral planning-scene snapshot -after the planning world has added its robots. This snapshot maps world robot IDs to -`RobotModelConfig` metadata so Viser can prepare current, target, and transient preview robot -visuals without `WorldMonitor` depending on Viser-specific hooks. Embedded Meshcat visualization -does not need extra setup because it observes the Drake world directly. - -When the Viser panel is enabled, it can call the existing manipulation execution path after a -fresh feasible plan is available and the current robot joints still match the plan start. +The Viser panel talks to the concrete `ManipulationOperator` bound into its +`VisualizationSession`. GUI callbacks enqueue operations through that operator +for target evaluation, planning, preview, execution, cancellation, reset, and +clear-plan actions. The panel owns only target drafts, selection state, and +callback generations; it does not touch `WorldSpec`, IK, planner objects, +`ManipulationModule`, `WorldMonitor`, or live Drake contexts directly. + +External manipulation visualizers are initialized from a backend-neutral +`VisualizationSession` after the planning world has added its robots. The +session contains static `PlanningSceneInfo` metadata: world robot IDs, +`RobotModelConfig` values, and resolved planning groups. Runtime joint state is +then pushed through `VisualizationStateFrame` updates so renderers do not poll +world/module state or own freshness policy. Embedded Meshcat visualization does +not need extra setup because it observes the Drake world directly. + +Previews use the stored synchronized `JointTrajectory` from the generated plan. +Viser projects the globally named trajectory into robot-local preview ghosts and +plays the stored timestamped points directly; optional preview duration only +scales the stored delays. Execute freshness is enforced by the manipulation +module/operator immediately before dispatch, not by Viser-side telemetry +snapshots. ### Perception + Agent diff --git a/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml b/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml new file mode 100644 index 0000000000..68b717479d --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-11 diff --git a/openspec/changes/clean-manipulation-visualization-boundary/design.md b/openspec/changes/clean-manipulation-visualization-boundary/design.md new file mode 100644 index 0000000000..7805a3e4a8 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/design.md @@ -0,0 +1,125 @@ +## Context + +Viser currently receives `ManipulationModule` and `WorldMonitor` directly. Its GUI and visualization adapter call a broad set of module and monitor APIs for topology, current state, FK, collision checks, IK, planning, preview, execution, cancellation, and plan freshness. The Viser layer also accepts local and global joint-name variants, composes complete robot states, projects generated paths for preview, and duplicates safety policy that belongs in manipulation. + +The visualization protocol reinforces this coupling: backends pull current state, receive `GeneratedPlan`, resolve planning groups at runtime, and separately show, animate, cancel, and hide preview ghosts. Execution independently turns the stored geometric path into trajectories, so preview and execution need not share one timing authority. + +The agreed domain language in `CONTEXT.md` distinguishes paths, synchronized trajectories, generated plans, target drafts, target evaluations, plan freshness, manipulation operators, and transient previews. This design applies those distinctions while preserving the existing visible Viser interaction contract. + +## Goals / Non-Goals + +**Goals:** + +- Materialize one synchronized trajectory immediately after geometric planning and store it with the path in `GeneratedPlan`. +- Make the stored trajectory the sole motion and timing authority for preview and execution. +- Let the existing `JointTrajectoryTask` execute a configured subset without commanding omitted joints. +- Replace Viser's direct module/monitor dependencies with one concrete, UI-neutral `ManipulationOperator`. +- Make visualization topology immutable after initialization and push current joint-state frames instead of allowing renderers to pull world state. +- Pass raw globally named trajectories to visualization backends and keep projection backend-local. +- Move execution freshness and authoritative action validation out of Viser. +- Remove obsolete adapters and representation conversions rather than retaining compatibility shims. + +**Non-Goals:** + +- Changing planner algorithms, IK algorithms, collision ownership, or world-backend kinematics. +- Providing transport-level simultaneous multi-robot state snapshots; planning keeps the current static-world assumption and existing state APIs. +- Supporting dynamic robot or planning-group topology after visualization initialization. +- Adding atomic multi-task coordinator dispatch, dynamic subset resource claims, or concurrent joint ownership semantics. +- Correcting the existing lack of coordinator task cancellation when manipulation execution is cancelled. +- Changing Viser controls, wording, layout, group selection behavior, target ghosts, or pose-gizmo interactions. +- Adding arbitrary pose-target frame transforms; operator pose requests initially require the world frame. + +## Decisions + +### 1. Materialize `GeneratedPlan` once after planner success + +`PlannerSpec.plan` continues to produce the planner's geometric `PlanningResult`. `ManipulationModule` validates the successful path, resolves selected global joint limits by exact name and requested order, calls one trajectory generator across all planned joints, validates the result, and constructs a complete `GeneratedPlan` containing both `path` and `trajectory`. + +Only complete generated plans are cached or exposed. Preview, execution, status, and completion timing MUST NOT regenerate or re-parameterize the path. The path remains geometric/diagnostic data; the trajectory is motion/timing authority. + +This keeps planner backends waypoint-focused and avoids a transient incomplete `GeneratedPlan` contract. + +### 2. Store only planned global joints and no separate baseline + +Both `GeneratedPlan.path` and `GeneratedPlan.trajectory` use canonical global joint names in selected planning-group order. `path[0]` and the first trajectory point represent the planned-joint start state. Joints outside the selected groups are not stored or commanded. + +No plan baseline is added. Under the accepted single-commander/static-state assumption, renderers overlay planned joints onto their latest current visual state. Execution sends only planned joints through the existing trajectory task. + +### 3. Extend the existing `JointTrajectoryTask` for partial trajectories + +The task validates that incoming trajectory names are a non-empty, unique subset of its configured joints and that every point matches the active width with finite values and valid timing. During execution it emits `JointCommandOutput` using the incoming active names, not the task's full configured list. + +`claim()` continues to claim the full configured joint set. This intentionally preserves current arbitration/preemption behavior while preventing commands to omitted joints. Dynamic claims are deferred. + +### 4. Initialize a static visualization session + +`PlanningSceneInfo` is extended with resolved `PlanningGroup` values in addition to robot configs keyed by world robot ID. `ManipulationModule` constructs one concrete `ManipulationOperator` after world setup and asks `WorldMonitor` to initialize the visualization with a session containing: + +- immutable scene/robot/planning-group metadata; and +- the optional concrete operator used by interactive backends. + +There is no `ManipulationOperatorSpec` or protocol because only one implementation exists. `VisualizationSpec` remains justified by its multiple backend implementations. + +### 5. Push fast state and keep slow operator status small + +`WorldMonitor` periodically builds a `VisualizationStateFrame` containing current joint states keyed by world robot ID and calls `VisualizationSpec.update_state(frame)`. The frame contains no staleness policy. Viser uses it for current URDF rendering, target initialization, current presets, and joint readout. + +`ManipulationOperator.status()` returns only slow-changing manipulation state, error, and stored-plan summary. It does not return topology or current joint telemetry. + +### 6. Use a raw trajectory visualization API + +The runtime visualization API is reduced to raw globally named trajectory animation and cancellation. Backends receive no `GeneratedPlan`, path, or runtime group IDs. They derive affected robots from trajectory joint names using static initialization metadata, preserve stored timestamps by default, and apply an explicit duration only as playback-rate scaling without mutating the trajectory. + +Preview ghosts are transient. Animation shows affected ghosts, plays the trajectory, and hides them in `finally`; cancellation invalidates playback and hides ghosts immediately. `show_preview` and `hide_preview` are removed. + +### 7. Use one concrete stateless manipulation operator facade + +The Viser panel receives only `ManipulationOperator` and its local scene/controller objects. The facade composes `ManipulationModule` and `WorldMonitor` and exposes synchronous typed operations: + +- compact `status`; +- joint/pose target evaluation; +- joint/pose planning; +- preview, execute, cancel, clear-plan, and reset. + +Target drafts, group selection, and callback generation IDs remain frontend-owned. Operator evaluation and planning calls are stateless with respect to drafts. + +Joint requests contain ordered group IDs and one exact globally named `JointState` whose names equal the selected groups' concatenated joints in group order. Pose requests contain explicit world-frame `PoseStamped` targets, optional auxiliary group IDs, and an optional globally named seed. + +Target evaluations are advisory and expose only selected global joints, group poses, and group diagnostics. Complete per-robot states used for FK/collision checks stay private. Planning consumes the original target request rather than a prior evaluation result. Planning returns a typed summary, never the complete `GeneratedPlan`. + +### 8. Keep authoritative freshness in manipulation + +Viser no longer polls `is_state_stale`, records complete robot snapshots, compares snapshots, or decides whether execution is fresh. Immediately before dispatch, `ManipulationModule.execute` resolves current planned joints and compares them against the first stored trajectory point within the configured tolerance. Missing, malformed, or mismatched state rejects execution. + +Omitted joints are intentionally ignored under the accepted single-commander assumption. + +### 9. Route preview cancellation through the existing outbound port + +`ManipulationOperator.cancel` calls `ManipulationModule.cancel`. The module routes preview cancellation through `WorldMonitor`, which calls `VisualizationSpec.cancel_preview_animation`; no layer knows whether Viser or another backend implements the port. + +The actual coordinator execution-cancellation gap remains explicitly out of scope. + +## Risks / Trade-offs + +- **Partial trajectories change task validation/output behavior** → Preserve full task claims, add strict subset/shape/timing validation, and retain full-width trajectory tests. +- **No stored baseline means omitted joints can differ between preview moments** → Accepted under the stated static/single-commander assumption; planned joints and timing remain identical. +- **Raw global trajectories make renderers aware of global joint naming** → Supply immutable robot/group metadata once and centralize backend-local name projection helpers. +- **A visualization callback can re-enter the same backend through operator preview** → Keep operator calls synchronous but run them from Viser's existing worker lane; retain visualization mutation/cancellation locking and generation tests. +- **Removing frontend freshness checks changes immediate button behavior** → Keep authoritative module rejection and return clear action errors; do not duplicate policy in state frames. +- **Direct cutover touches several layers simultaneously** → Land model/materialization and task subset support first, then protocol/session/operator migration, then delete legacy methods and adapters after all backends compile. +- **Exact API names may need adjustment against construction code** → Preserve the directional contracts and responsibilities even if final concrete class or method names change during implementation. + +## Migration Plan + +1. Add trajectory materialization and partial-task support with focused tests while preserving existing visualization calls temporarily inside the implementation branch. +2. Add static visualization session/state-frame DTOs and the concrete operator; initialize them after world setup. +3. Migrate WorldMonitor, Drake, and Viser to `initialize`, `update_state`, raw `animate_trajectory`, and cancellation. +4. Migrate the Viser panel to typed operator requests/results and pushed current state; remove direct module/monitor access, snapshot freshness logic, and `GroupPanelBackend`. +5. Remove old visualization methods and all lazy path-to-trajectory/preview projection paths. +6. Run full manipulation, planning, visualization, control, lifecycle, and multi-robot tests plus mypy, Ruff, diff, and strict OpenSpec validation. + +Rollback is a source revert because this is an internal breaking boundary with no persisted data migration. + +## Open Questions + +- Final concrete DTO and method names are provisional until implementation verifies construction and lifecycle call sites; their responsibilities and directionality are fixed by this design. diff --git a/openspec/changes/clean-manipulation-visualization-boundary/proposal.md b/openspec/changes/clean-manipulation-visualization-boundary/proposal.md new file mode 100644 index 0000000000..847ca42ac8 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/proposal.md @@ -0,0 +1,33 @@ +## Why + +The manipulation visualization boundary currently lets Viser pull from both `ManipulationModule` and `WorldMonitor`, so UI code performs planning-group lookup, state normalization, FK/collision composition, preview projection, and execution freshness checks. Preview and execution also derive different representations from an untimed stored path, making the boundary difficult to reason about and easy to drift. + +## What Changes + +- **BREAKING** Extend `GeneratedPlan` to contain both its geometric path and one synchronized, globally named `JointTrajectory`; parameterize the path exactly once immediately after `PlannerSpec.plan` succeeds and before the plan is cached or exposed. +- Make the stored trajectory the sole motion/timing authority for preview and execution. Preserve the geometric path for inspection and diagnostics. +- **BREAKING** Replace the pull-oriented visualization API with one-shot static initialization, pushed current-state frames, raw synchronized-trajectory animation, and cancellation. Remove `show_preview`, `hide_preview`, and `animate_plan(GeneratedPlan)`. +- Include resolved planning groups in immutable visualization initialization metadata. Visualization topology does not change during a session. +- Add one concrete `ManipulationOperator` facade that composes `ManipulationModule` and `WorldMonitor` for typed target evaluation and manipulation actions. Viser receives this facade through visualization initialization instead of receiving the raw module and monitor. +- Keep target drafts and asynchronous request sequencing frontend-owned. Make target evaluation advisory and require planning to consume the original complete target request. +- Move plan freshness enforcement into `ManipulationModule.execute`, comparing current planned joints with the stored trajectory start. Remove Viser-side robot snapshots, stale-state policy, and execution safety decisions. +- Push only current joint states to visualizations at visualization rate; freshness remains an authoritative manipulation concern and does not cross the rendering API. +- Update the existing `JointTrajectoryTask` to accept a non-empty configured joint subset while retaining its full configured resource claim. It emits commands only for trajectory joints, so omitted joints are not commanded. +- Preserve the existing Viser controls and interaction contract while making preview ghosts transient: completion or cancellation hides them immediately. + +## Capabilities + +### New Capabilities +- `manipulation-visualization-boundary`: Static visualization sessions, pushed joint-state frames, raw synchronized-trajectory preview, cancellation, and backend isolation from manipulation internals. +- `manipulation-operator`: The concrete UI-neutral facade, canonical target requests/evaluations, compact status, and action semantics used by interactive visualization panels. +- `partial-joint-trajectory-execution`: Validation and execution of configured joint subsets through the existing `JointTrajectoryTask` while preserving full resource claims. + +### Modified Capabilities +- `manipulation-module-group-api`: Generated plans become materialized path-plus-trajectory artifacts; preview and execution consume the stored trajectory and module-owned execution freshness. + +## Impact + +- Affects manipulation planning models and module materialization/execution, visualization protocols and scene metadata, `WorldMonitor`, Drake and Viser backends, Viser panel/backend/state code, trajectory control task validation/output, tests, and manipulation documentation. +- Removes direct `ManipulationModule` and `WorldMonitor` dependencies from Viser GUI/renderer construction. +- Replaces existing visualization method signatures and `GeneratedPlan` construction sites; this is an intentional internal breaking migration without compatibility shims. +- Does not change planner algorithms, robot/world collision ownership, coordinator batching, actual hardware execution cancellation, or the visible Viser interaction design. diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md new file mode 100644 index 0000000000..58453778e5 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md @@ -0,0 +1,56 @@ +## ADDED Requirements + +### Requirement: Successful planning must materialize one complete generated plan +After `PlannerSpec.plan` returns a successful geometric path, ManipulationModule MUST validate the selected global waypoints and parameterize them exactly once into one synchronized globally named `JointTrajectory`. It MUST cache or expose a `GeneratedPlan` only when both path and trajectory are complete and mutually consistent. + +#### Scenario: Selected-group planning succeeds +- **WHEN** a planner returns valid waypoints for the requested planning groups +- **THEN** ManipulationModule stores one generated plan containing the original path and one trajectory over the same selected global joints + +#### Scenario: Parameterization fails +- **WHEN** joint limits, waypoint data, or generated trajectory validation fails +- **THEN** planning fails and no generated plan is cached or exposed + +#### Scenario: Multiple robots are selected +- **WHEN** selected group joints span multiple robots +- **THEN** one generator call creates one shared relative trajectory clock across all planned joints + +### Requirement: Stored trajectory must be preview and execution authority +Preview, execution, completion timing, and status logic MUST consume the stored generated-plan trajectory without re-projecting or re-parameterizing the geometric path. + +#### Scenario: Plan is previewed and executed +- **WHEN** a stored generated plan is previewed and later executed +- **THEN** both operations consume the same trajectory points, joint order, and timestamps + +#### Scenario: Preview duration is overridden +- **WHEN** a caller requests a different display duration +- **THEN** only visualization playback rate changes and the stored trajectory remains unchanged + +### Requirement: Generated trajectories must contain only planned joints +The generated path and trajectory MUST contain exactly the selected planning groups' global joint names in selection order. Joints outside the selected groups MUST NOT be added to the generated plan or execution command. + +#### Scenario: Planning group omits a robot joint +- **WHEN** a selected group excludes another configured robot joint +- **THEN** the excluded joint appears in neither the generated trajectory nor the trajectory task command + +### Requirement: Manipulation module must own execution freshness +Immediately before dispatch, ManipulationModule MUST compare current planned-joint positions with the stored trajectory's first point within configured tolerance. Missing, malformed, duplicate, reordered, stale, or mismatched planned-joint state MUST reject execution before task dispatch. + +#### Scenario: Planned joints still match +- **WHEN** all current planned joints match the trajectory start within tolerance +- **THEN** execution may translate and dispatch the stored trajectory + +#### Scenario: Planned joint moved after planning +- **WHEN** any current planned joint differs from the trajectory start beyond tolerance +- **THEN** execution is rejected without relying on frontend snapshots + +### Requirement: Execution must translate only at the coordinator boundary +Execution MUST split the stored global trajectory by affected robot without changing timing, translate selected local names to coordinator names at dispatch, and submit the resulting partial trajectory to the existing configured trajectory task. + +#### Scenario: Selected joints use coordinator aliases +- **WHEN** a robot config maps selected local joint names to coordinator names +- **THEN** dispatch preserves stored points and timestamps while translating only those selected names + +#### Scenario: Plan spans multiple robots +- **WHEN** planned global joints span multiple configured trajectory tasks +- **THEN** each task receives its selected local subset on the same stored relative clock using the existing sequential submission behavior diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md new file mode 100644 index 0000000000..5ea9e4f106 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md @@ -0,0 +1,73 @@ +## ADDED Requirements + +### Requirement: Interactive manipulation must use one concrete operator facade +The system MUST provide one concrete UI-neutral manipulation operator that composes `ManipulationModule` and `WorldMonitor`. It MUST NOT introduce an operator Spec or Protocol when no alternate implementation exists. + +#### Scenario: Visualization session is created +- **WHEN** ManipulationModule completes world initialization +- **THEN** it constructs one operator and binds that concrete instance into the visualization session + +### Requirement: Operator dynamic reads must remain compact +The operator MUST expose slow-changing manipulation state, error, and stored-plan summary through one compact status read. Static topology and high-rate joint states MUST NOT be returned by operator status. + +#### Scenario: Panel refreshes status +- **WHEN** Viser refreshes module and plan status +- **THEN** it reads compact operator status without copying robot configs or joint telemetry + +### Requirement: Joint target requests must use canonical global joints +Joint target requests MUST contain ordered planning-group IDs and one globally named `JointState` whose names exactly equal the selected groups' concatenated joint names in group order. Local aliases, suffix matching, missing joints, duplicate joints, and extraneous joints MUST be rejected. + +#### Scenario: Multiple groups submit a joint draft +- **WHEN** a frontend submits targets for ordered groups across one or more robots +- **THEN** the operator validates one flat global joint state in that exact selection order + +### Requirement: Pose target requests must carry explicit frames +Pose target requests MUST contain `PoseStamped` values keyed by pose-capable planning-group ID, optional ordered auxiliary group IDs, and an optional globally named seed. Initial operator support MUST require every pose target to use the world frame. + +#### Scenario: World-frame pose target is evaluated +- **WHEN** a frontend submits a world-frame target for a pose-capable group +- **THEN** the operator evaluates it without inventing or rewriting its frame + +#### Scenario: Unsupported pose frame is submitted +- **WHEN** a pose target uses a frame other than world +- **THEN** the operator rejects the request with a clear validation result + +### Requirement: Target evaluations must be advisory and selected-domain only +Target evaluation MUST return feasibility, status, message, selected globally named joints when available, group poses, and per-group diagnostics. Complete per-robot states used for FK or collision checks MUST remain internal, and an evaluation MUST NOT authorize planning or execution. + +#### Scenario: Joint target is evaluated +- **WHEN** the operator composes complete robot states for collision and FK checks +- **THEN** the frontend receives only selected-domain results and diagnostics + +#### Scenario: Pose target is evaluated +- **WHEN** IK produces selected joints +- **THEN** the operator returns those selected global joints and group results without exposing internal complete states + +### Requirement: Planning must consume original complete target requests +Operator planning methods MUST consume the original joint or pose target request and MUST NOT accept a prior advisory evaluation as planning authority. Planning MUST return a typed action summary rather than the complete `GeneratedPlan`. + +#### Scenario: Target changes after evaluation +- **WHEN** a frontend plans a newer target draft after an older evaluation completed +- **THEN** planning validates and plans the submitted request independently of the older evaluation + +#### Scenario: Planning succeeds +- **WHEN** manipulation caches a complete generated plan +- **THEN** the operator returns success, message, group IDs, waypoint count, and duration without returning path or trajectory data + +### Requirement: Target drafts and scheduling must remain frontend-owned +The operator MUST be synchronous and stateless with respect to editable target drafts, group selection, request generations, and frontend worker scheduling. + +#### Scenario: Evaluation results arrive out of order +- **WHEN** Viser runs synchronous operator calls in workers and a newer target generation exists +- **THEN** Viser discards the stale result using its own generation identity + +### Requirement: Operator actions must preserve domain ownership +Preview, execute, cancel, clear-plan, and reset methods MUST delegate to manipulation behavior and return typed action results. Cancellation MUST route preview cancellation through WorldMonitor's configured visualization without directly addressing Viser. + +#### Scenario: Operator cancels during preview +- **WHEN** the frontend calls operator cancel +- **THEN** ManipulationModule cancellation routes through WorldMonitor to visualization cancellation + +#### Scenario: Operator clears a plan +- **WHEN** the frontend calls clear-plan +- **THEN** the stored generated plan is discarded without introducing persistent preview state diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md new file mode 100644 index 0000000000..0f7b46b791 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md @@ -0,0 +1,64 @@ +## ADDED Requirements + +### Requirement: Visualization initialization must receive immutable topology +The visualization protocol MUST initialize each backend once with robot configurations, resolved planning groups, and an optional concrete manipulation operator. Robot and planning-group topology MUST remain unchanged for the initialized session. + +#### Scenario: Interactive visualization initializes +- **WHEN** manipulation world setup completes with an interactive visualization backend +- **THEN** the backend receives all robot IDs, robot configs, resolved planning groups, and the concrete manipulation operator in one initialization session + +#### Scenario: Non-interactive visualization initializes +- **WHEN** a rendering backend does not provide an operator panel +- **THEN** it initializes from the same static topology and ignores the optional operator + +### Requirement: Visualization state must be pushed without world access +WorldMonitor MUST push current joint-state frames to visualization backends at visualization update cadence. A visualization backend MUST NOT pull current state, freshness, FK, collision, planning, or execution data from `WorldMonitor` or `ManipulationModule`. + +#### Scenario: Current robot state changes +- **WHEN** WorldMonitor publishes a visualization update +- **THEN** the backend receives current joint states keyed by initialized world robot IDs + +#### Scenario: State freshness changes +- **WHEN** telemetry becomes stale +- **THEN** no freshness policy is added to the visualization frame and authoritative manipulation actions validate state independently + +### Requirement: Visualization preview must consume raw synchronized trajectories +The visualization protocol MUST animate the stored globally named `JointTrajectory` directly and MUST NOT receive `GeneratedPlan` or regenerate timing from a geometric path. + +#### Scenario: Stored trajectory is previewed +- **WHEN** manipulation previews a generated plan without a display-duration override +- **THEN** the backend plays the raw trajectory using its stored timestamps + +#### Scenario: Display duration is overridden +- **WHEN** preview specifies a display duration +- **THEN** the backend scales playback delays without mutating or re-parameterizing the stored trajectory + +#### Scenario: Multiple robots are planned +- **WHEN** trajectory joint names span multiple initialized robots +- **THEN** the backend projects those joints through static topology and advances every affected preview ghost on the same stored clock + +### Requirement: Visualization preview must be transient +The visualization protocol MUST expose only trajectory animation and cancellation for preview lifecycle. Preview ghosts MUST disappear after completion, cancellation, error, replacement, or backend close. + +#### Scenario: Preview completes +- **WHEN** the final trajectory time is displayed +- **THEN** all preview ghosts are hidden immediately + +#### Scenario: Preview is cancelled +- **WHEN** cancellation occurs during playback +- **THEN** stale frames stop mutating the scene and all preview ghosts are hidden immediately + +#### Scenario: Preview is replaced +- **WHEN** a new trajectory begins while another preview is active +- **THEN** the old generation is cancelled and cannot mutate the new preview + +### Requirement: Visualization backends must remain domain-isolated +Viser GUI and renderer construction MUST NOT receive raw `ManipulationModule` or `WorldMonitor` dependencies. Runtime render methods MUST NOT perform planning-result validation, target evaluation, execution freshness checks, or action dispatch outside the concrete manipulation operator. + +#### Scenario: Viser panel starts +- **WHEN** the visualization session includes an operator +- **THEN** the panel uses only that operator for domain reads and actions + +#### Scenario: Preview renders +- **WHEN** Viser receives a valid raw trajectory +- **THEN** it performs only topology/name projection, visual-state overlay, playback, and ghost cleanup diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md new file mode 100644 index 0000000000..4f27ffa242 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md @@ -0,0 +1,45 @@ +## ADDED Requirements + +### Requirement: Existing trajectory tasks must accept configured joint subsets +`JointTrajectoryTask` MUST accept a trajectory only when its joint names are a non-empty, unique subset of the task's configured joints and every trajectory point matches that active joint width. Unknown, duplicate, empty, malformed, non-finite, or invalidly timed input MUST be rejected before execution state changes. + +#### Scenario: Selected planning group trajectory is submitted +- **WHEN** a task configured for a complete arm receives a valid trajectory containing only that arm's selected planning-group joints +- **THEN** the existing task accepts and executes the trajectory + +#### Scenario: Unknown joint is submitted +- **WHEN** a trajectory contains a joint outside the task configuration +- **THEN** the task rejects it without entering executing state + +#### Scenario: Point width is malformed +- **WHEN** a trajectory point's positions or velocities do not match the trajectory joint-name count +- **THEN** the task rejects it before storing the trajectory + +### Requirement: Partial trajectory output must command only active joints +While executing a partial trajectory, `JointTrajectoryTask` MUST emit `JointCommandOutput` names and values only for the trajectory's active joint subset. It MUST NOT fill or command configured joints omitted from the trajectory. + +#### Scenario: Partial trajectory is sampled +- **WHEN** a two-joint trajectory executes on a six-joint configured task +- **THEN** each output contains exactly those two joint names and two sampled positions + +#### Scenario: Full trajectory is sampled +- **WHEN** a trajectory contains every configured task joint +- **THEN** existing full-width output behavior remains unchanged + +### Requirement: Resource claims must remain configuration-wide +`JointTrajectoryTask.claim()` MUST continue to claim the task's full configured joint set regardless of the active trajectory subset. + +#### Scenario: Partial trajectory executes +- **WHEN** the active trajectory omits configured joints +- **THEN** arbitration and preemption still use the task's full configured resource claim + +### Requirement: Partial trajectory lifecycle must clear active names safely +Replacement, completion, cancellation, reset, and fault paths MUST keep active trajectory joint names consistent with the stored trajectory and MUST NOT leak names from a prior execution. + +#### Scenario: Partial trajectory is replaced +- **WHEN** an active subset trajectory is replaced by a different valid subset +- **THEN** subsequent outputs use only the replacement trajectory's names + +#### Scenario: Task resets +- **WHEN** a completed or aborted task resets +- **THEN** its stored trajectory and active subset are cleared while its configured claim remains unchanged diff --git a/openspec/changes/clean-manipulation-visualization-boundary/tasks.md b/openspec/changes/clean-manipulation-visualization-boundary/tasks.md new file mode 100644 index 0000000000..323099d4e8 --- /dev/null +++ b/openspec/changes/clean-manipulation-visualization-boundary/tasks.md @@ -0,0 +1,48 @@ +## 1. Materialize Generated Plans + +- [x] 1.1 Extend `GeneratedPlan` with a required synchronized `JointTrajectory`, migrate construction sites, and add model/import/round-trip tests proving path and trajectory remain distinct representations. +- [x] 1.2 Add one post-planner materialization path in `ManipulationModule` that validates selected global waypoints, resolves velocity/acceleration limits by exact global name and selection order, and invokes one trajectory generator across all planned joints. +- [x] 1.3 Validate generated trajectory names, dimensions, finite values, timestamps, positive duration, and ordered waypoint boundaries before caching; ensure planning/parameterization failure exposes no generated plan. +- [x] 1.4 Remove non-dispatch lazy path-to-trajectory generation from status/completion timing and isolate the remaining dispatch-only bridge for Phase 2 removal after partial trajectory task support lands. +- [x] 1.5 Add materialization tests for reordered groups, heterogeneous multi-robot limits, malformed/non-finite waypoints, invalid limits, generator failure, exactly one generation call, shared timing, and zero generation after caching. + +## 2. Execute Partial Joint Trajectories + +- [x] 2.1 Update the existing `JointTrajectoryTask.execute` validation to accept only non-empty unique configured subsets with valid point widths, finite values, and strictly increasing timing while preserving full-width behavior. +- [x] 2.2 Make task computation emit only the active trajectory joint names and sampled values while leaving `claim()` on the full configured joint set and clearing active subset state on replacement, completion, cancellation, fault, and reset. +- [x] 2.3 Update manipulation execution to split the stored global trajectory by robot, preserve every timestamp/value, translate selected names only at coordinator dispatch, and never fill or command omitted joints. +- [x] 2.4 Add control/manipulation tests for valid subsets, unknown/duplicate/empty names, malformed positions/velocities/timing, full-width compatibility, full configured claims, replacement/reset lifecycle, coordinator translation, omitted-joint non-commanding, and multi-robot shared clocks. + +## 3. Introduce the Visualization Session Boundary + +- [x] 3.1 Extend `PlanningSceneInfo` with resolved immutable planning groups and add visualization session/state-frame DTOs containing static topology, the optional concrete operator, and pushed current joint states respectively. +- [x] 3.2 Replace `VisualizationSpec` pull/group/plan methods with one-shot initialization, `update_state`, raw globally named `animate_trajectory`, cancellation, URL, and close; add structural/import tests for the complete protocol. +- [x] 3.3 Make `WorldMonitor` initialize visualization once, push current joint-state frames at visualization cadence, serialize renderer mutations, forward raw stored trajectories, and route cancellation without exposing freshness policy. +- [x] 3.4 Migrate Drake/Meshcat visualization to static topology and raw trajectory playback, preserving stored non-uniform timestamps, optional playback scaling, synchronized multi-robot ticks, generation cancellation, and automatic ghost cleanup. +- [x] 3.5 Add WorldMonitor and Drake tests for immutable group metadata, pushed state frames, no state pull, raw timing/scaling, trajectory immutability, multi-robot shared ticks, replacement/cancel/error/close cleanup, and removal of legacy visualization methods. + +## 4. Add the Concrete Manipulation Operator + +- [x] 4.1 Add concrete operator status, canonical joint/pose request, target evaluation, plan summary, action result, and visualization session types without introducing a new operator Protocol or Spec. +- [x] 4.2 Implement exact ordered global-joint request validation and explicit world-frame pose request validation, rejecting aliases, duplicates, omissions, extras, unsupported groups/frames, and malformed seeds. +- [x] 4.3 Move joint/pose target composition, IK, FK, per-robot state validity, and selected-domain evaluation results from Viser-local helpers into the operator while keeping complete robot states private. +- [x] 4.4 Implement synchronous stateless planning methods that consume original requests, delegate to manipulation APIs, and return summaries without exposing `GeneratedPlan`; implement compact slow-changing status and action methods. +- [x] 4.5 Move execute freshness enforcement into `ManipulationModule` by comparing current planned global joints with the stored trajectory start immediately before dispatch; remove dependence on frontend robot snapshots and stale-state gates. +- [x] 4.6 Construct one operator after module/world setup, bind it into visualization initialization, preserve indirect module-to-world-to-visualization cancellation, and add lifecycle/action/freshness tests. + +## 5. Migrate Viser to the Clean Boundary + +- [x] 5.1 Initialize Viser from static robot/group metadata and build deterministic global-joint-to-robot/local-joint maps without retaining `ManipulationModule` or `WorldMonitor` dependencies. +- [x] 5.2 Consume pushed state frames for current URDFs, joint readout, target initialization, and Current presets; remove direct state/freshness polling and Viser-side stale policy. +- [x] 5.3 Replace generated-plan preview projection and preview-track resampling with raw stored-trajectory projection/playback; support duration scaling only, hide ghosts on completion/cancel/error/replacement/close, and prevent stale frames from mutating scene handles. +- [x] 5.4 Migrate the panel to the concrete operator's typed status, target evaluation, planning, preview, execution, cancellation, clear, and reset calls while keeping target drafts and callback generations frontend-owned. +- [x] 5.5 Remove `GroupPanelBackend`, direct module/monitor calls, local/global alias normalization, complete-state composition, robot snapshot matching, and obsolete preview show/hide state from Viser code. +- [x] 5.6 Preserve the normative Viser UI controls and interactions exactly and add granular tests for group selection, presets, sliders, pose gizmos, advisory evaluation, action results, pushed current state, raw preview timing, auto-hide, cancellation races, and absence of forbidden dependencies. + +## 6. Cleanup, Documentation, and Validation + +- [x] 6.1 Remove all production/test references to `animate_plan`, `show_preview`, `hide_preview`, lazy path projection/parameterization, Viser `ManipulationModule`/`WorldMonitor` access, and legacy panel adapter APIs. +- [x] 6.2 Update manipulation/planning-group/visualization documentation and `CONTEXT.md` to describe materialized generated plans, partial planned-joint commands, the concrete operator, pushed visualization state, and transient previews. +- [x] 6.3 Run focused planner/materialization, manipulation module/unit/reservation, WorldMonitor, Drake, Viser/factory/lifecycle, control task/coordinator, and multi-robot end-to-end tests with no relevant skips. +- [x] 6.4 Run targeted mypy over every changed production file, Ruff check/format, `git diff --check`, and `openspec validate clean-manipulation-visualization-boundary --strict`; audit the final diff for unrelated planner/control/UI changes. +- [x] 6.5 Complete a final architecture and safety review confirming one trajectory-generation call, stored-timing preview/execution parity, partial-output/full-claim semantics, module-owned freshness, renderer isolation, exact Viser interaction parity, and no compatibility shims. diff --git a/openspec/changes/viser-group-planning-ui/design.md b/openspec/changes/viser-group-planning-ui/design.md index 0e201b7004..089ebc6a17 100644 --- a/openspec/changes/viser-group-planning-ui/design.md +++ b/openspec/changes/viser-group-planning-ui/design.md @@ -21,8 +21,20 @@ Viser must let users select planning groups, move joint and pose targets, see fe - Treat planning-group IDs as the UI's stable selection keys. - Validate base-link/root assumptions before rendering URDFs under base poses. - Keep manual demo/checklist coverage for UI feel in addition to unit tests. +- Treat `cc/spec/movegroup@0edb8d3dd` as the sole source for visible controls and interaction behavior; do not derive behavior from the upstream extraction commit. +- Adapt the reference UI to PR 4's existing explicit group APIs rather than importing manipulation or planner changes from the reference branch. +- Keep execution trajectory parameterization in the manipulation/module layer. Group-native preview renderers may map canonical `GeneratedPlan` waypoints onto the requested display duration, but must not reinterpret them as execution timing. +- Use the reference group-native visualization contract: `show_preview(Sequence[PlanningGroupID])`, `hide_preview(Sequence[PlanningGroupID])`, and `animate_plan(GeneratedPlan, duration: float = 3.0)`. `ManipulationModule.preview_plan` supplies `1.0` when duration is omitted, matching the normative reference. Update Viser and the existing Meshcat implementation together; do not import unrelated source-branch planning changes. +- Treat preview as one validated transaction. Resolve and deduplicate affected robots, snapshot each robot's non-selected joints once, project selected global joints into complete local frames from that fixed baseline, reject malformed/unknown/incomplete input before showing ghosts, then update every ghost once per shared tick before sleeping once. Unequal renderer frame counts map onto the same normalized clock. +- Use an animation generation/cancellation token. Cancel, clear, close, or a replacement preview invalidates the old generation so it cannot mutate removed or replacement scene handles. Serialize preview-scene mutation against periodic visualization publication. +- Group selection owns a monotonically increasing epoch. Target/plan callbacks commit only if both their operation sequence and selection epoch still match. Execute snapshots every affected robot immediately before dispatch, compares every selected group's ordered joints, and calls `ManipulationModule.execute()` without a robot filter. +- Execute snapshots every affected robot immediately before dispatch and delegates the stored plan without a robot filter. An external caller can unavoidably replace that stored plan between the panel freshness check and dispatch. +- Do not copy absent reference APIs such as `ManipulationModule.forward_kinematics` or `check_collision`. Use PR4 `WorldMonitor.get_group_ee_pose` for group FK and `WorldMonitor.is_state_valid` on composed full per-robot target states. Multi-robot validity is a conjunction of per-robot checks, not proof of simultaneous target-target collision freedom; pose IK retains PR4's coordinated collision check. ## Risks / Trade-offs - UI tests can pass while interaction feels wrong; require a small manual checklist before review. - This PR is still large, but isolating it keeps visual review focused. +- A whole-commit cherry-pick would leak unrelated planner and module changes; extraction must be semantic and visualization-scoped. +- The preview contract is the sole approved scope expansion outside visualization. Keep its module/monitor/protocol/world changes minimal and independently testable. +- A compatibility `robot_name` argument may validate that the named robot is affected, but MUST NOT trim a multi-robot generated plan before preview. diff --git a/openspec/changes/viser-group-planning-ui/proposal.md b/openspec/changes/viser-group-planning-ui/proposal.md index 54137b817c..99079912d3 100644 --- a/openspec/changes/viser-group-planning-ui/proposal.md +++ b/openspec/changes/viser-group-planning-ui/proposal.md @@ -20,6 +20,9 @@ The Viser changes are large and have different review criteria from planning alg ## Impact - Base branch: PR 4 `manipulation-module-group-api`. -- Reference implementation: `cc/spec/movegroup`. +- Normative UI reference: `cc/spec/movegroup@0edb8d3dd`. User-visible deviations require explicit approval. +- API base: PR 4 at `cc/planning_group/main@05c25787a`; its explicit planning-group API scope remains unchanged. +- The upstream extraction commit `origin/cc/planning_group/viser@737fb3381` is not an implementation source. - Primary files: `dimos/manipulation/visualization/viser/*`, `dimos/manipulation/visualization/types.py`, `dimos/manipulation/visualization/test_factory.py`. -- Out of scope: planning model/backend/algorithm changes except as already supplied by earlier PRs. +- Supporting preview seam: group-native visualization protocol, world-monitor forwarding, manipulation-module preview routing, and the existing Meshcat implementation. +- Out of scope: all other manipulation-module, planning model/backend/algorithm, and control changes except as already supplied by earlier PRs. diff --git a/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md b/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md index a653991fcd..f28750806f 100644 --- a/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md +++ b/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md @@ -7,6 +7,77 @@ The Viser planning panel MUST present and use planning-group selections for grou - **WHEN** a user selects a planning group in the panel - **THEN** joint controls, pose target controls, and preview state apply to that group +### Requirement: Viser panel must preserve the normative reference interaction contract +The Viser panel MUST match the controls, labels, ordering, defaults, selection behavior, presets, joint sliders, pose gizmos, target ghosts, feasibility colors, preview behavior, execute gate, and clear behavior from `cc/spec/movegroup@0edb8d3dd`. + +#### Scenario: Extracted panel is reviewed +- **WHEN** the extracted panel is compared with the normative reference +- **THEN** no user-visible deviation is present unless separately approved + +### Requirement: Viser must preserve the cleaned planning-group API boundary +The Viser implementation MUST consume PR 4's explicit planning-group APIs without adding planner APIs. The only manipulation/visualization API change permitted is the group-native preview transaction required below. + +#### Scenario: Backend adaptation is required +- **WHEN** reference UI code expects a different backend shape +- **THEN** the visualization backend translates the call without changing planner algorithms or control production code + +### Requirement: Preview must be group-native and synchronized +The visualization protocol MUST expose `show_preview` and `hide_preview` over ordered planning-group IDs and `animate_plan` over one generated plan plus a total display duration. It MUST preview the plan as one validated transaction so all affected robots advance on a shared animation clock. + +#### Scenario: Selected groups span multiple robots +- **WHEN** the user previews a fresh multi-robot plan +- **THEN** all affected preview ghosts advance together from the same generated plan + +#### Scenario: Preview duration is omitted +- **WHEN** `ManipulationModule.preview_plan` is called without a duration +- **THEN** it sends the full generated plan to `animate_plan` with a total display duration of `1.0` second + +#### Scenario: Visualization is called directly +- **WHEN** `animate_plan` is called without an explicit duration +- **THEN** its total display duration defaults to `3.0` seconds + +#### Scenario: A plan controls a subset of robot joints +- **WHEN** a generated plan contains selected global joints +- **THEN** each affected robot's unselected joints come from one fixed pre-playback snapshot for every animation frame + +#### Scenario: Preview input is invalid +- **WHEN** a plan contains unknown groups, malformed joint states, or cannot form complete affected-robot frames +- **THEN** preview rejects the transaction before showing any ghost + +#### Scenario: Preview is replaced or stopped +- **WHEN** cancel, clear, close, or a new preview invalidates an active preview generation +- **THEN** old animation frames no longer mutate scene handles + +### Requirement: Panel freshness must cover the full selected target set +The panel MUST associate asynchronous results with a selection epoch and MUST snapshot ordered joints for every affected robot when a plan becomes executable. + +#### Scenario: Selection changes during evaluation +- **WHEN** a planning group is toggled while target evaluation is in flight +- **THEN** the stale result does not update target state or visuals + +#### Scenario: A newer operation supersedes an older operation +- **WHEN** callbacks arrive out of order for the same selection epoch +- **THEN** only the callback whose operation sequence is current may update panel or scene state + +#### Scenario: Execute spans multiple robots +- **WHEN** every selected group's ordered joints still match the stored all-robot snapshot +- **THEN** execute delegates the full stored plan without applying a robot filter + +#### Scenario: Another caller replaces the stored module plan +- **WHEN** replacement occurs between the panel freshness check and execution +- **THEN** the external-plan replacement race is unavoidable + +### Requirement: Joint-target feasibility must use existing PR 4 world APIs +The visualization backend MUST compose full per-robot target states and call `WorldMonitor.is_state_valid`; it MUST NOT add or copy manipulation-module FK or collision APIs. + +#### Scenario: Multiple robots have joint targets +- **WHEN** each composed per-robot state is valid +- **THEN** the panel reports the conjunction of those checks without claiming simultaneous target-target collision proof + +#### Scenario: Pose targets are evaluated +- **WHEN** selected pose groups and auxiliary groups form a target set +- **THEN** the backend uses PR 4 `inverse_kinematics` with coordinated collision checking, then validates each affected robot's composed full target state through `WorldMonitor.is_state_valid` + ### Requirement: Viser preview must reflect group feasibility and current state The Viser UI MUST show target feasibility and preview validity for the selected planning group. @@ -19,4 +90,4 @@ The Viser panel MUST prevent execution when no fresh plan matches the current se #### Scenario: Current state no longer matches preview - **WHEN** the robot state changes after planning -- **THEN** the execute action is rejected until a fresh matching plan is generated +- **THEN** the panel rejects execute until all selected-group snapshots match again diff --git a/openspec/changes/viser-group-planning-ui/tasks.md b/openspec/changes/viser-group-planning-ui/tasks.md index 0e9523358e..b9a75ae999 100644 --- a/openspec/changes/viser-group-planning-ui/tasks.md +++ b/openspec/changes/viser-group-planning-ui/tasks.md @@ -1,17 +1,17 @@ ## 1. Extract Viser UI changes -- [ ] 1.1 Start from PR 4 and use `cc/spec/movegroup` as reference. -- [ ] 1.2 Extract Viser panel/backend/scene/state/visualizer group-aware changes. -- [ ] 1.3 Remove or replace the old adapter code only as required by this UI slice. -- [ ] 1.4 Extract visualization type/factory updates needed by Viser. +- [x] 1.1 Start from PR 4 and use only `cc/spec/movegroup@0edb8d3dd` as the normative UI reference; do not rely on the upstream extraction commit. +- [x] 1.2 Extract Viser panel/backend/scene/state/visualizer group-aware changes and the minimal group-native shared-clock preview seam. +- [x] 1.3 Remove or replace the old adapter code only as required by this UI slice. +- [x] 1.4 Extract visualization type/factory updates needed by Viser. ## 2. Tests and manual check -- [ ] 2.1 Bring over Viser unit tests and lifecycle tests. -- [ ] 2.2 Run the manual checklist: group selector, joint sliders, pose gizmo, target ghost, infeasible target color, path preview, execute gate, clear path. +- [x] 2.1 Bring over Viser unit/lifecycle tests and add protocol, monitor forwarding, module routing, Meshcat projection, shared-tick, all-robot freshness, and preview cancellation race coverage. +- [x] 2.2 Run the manual checklist: group selector, joint sliders, pose gizmo, target ghost, infeasible target color, path preview, execute gate, clear path. ## 3. Validation -- [ ] 3.1 Run `uv run pytest dimos/manipulation/visualization/test_factory.py dimos/manipulation/visualization/viser/test_*.py -q`. -- [ ] 3.2 Run targeted mypy on changed Viser production files if practical. -- [ ] 3.3 Verify no planner/backend implementation changes are included beyond compile fixes against PR 4. +- [x] 3.1 Run `uv run pytest dimos/manipulation/test_manipulation_unit.py dimos/manipulation/planning/monitor/test_world_monitor.py dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/visualization/test_factory.py dimos/manipulation/visualization/viser/test_*.py -q` and `openspec validate viser-group-planning-ui --strict`. +- [x] 3.2 Run mandatory targeted mypy on changed Viser production files plus `manipulation_module.py`, `protocols.py`, `world_monitor.py`, and `drake_world.py`. +- [x] 3.3 Verify changes outside visualization are limited to the approved group-native preview protocol/forwarding/implementations, with no planner/backend algorithm or control changes. diff --git a/pyproject.toml b/pyproject.toml index b9cc00d79f..8251d7b923 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -607,6 +607,8 @@ ignore = [ "dimos/dashboard/dimos.rbl", "dimos/web/dimos_interface/themes.json", "dimos/manipulation/manipulation_module.py", + "dimos/manipulation/planning/world/drake_world.py", + "dimos/manipulation/visualization/viser/gui.py", "dimos/manipulation/visualization/viser/test_viser_visualization.py", "dimos/navigation/cmu_nav/modules/*/main.cpp", "dimos/navigation/cmu_nav/common/*.hpp", From 78ea1d6dae7806c000df7060d90e85c22fdb094e Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 14 Jul 2026 17:06:13 -0700 Subject: [PATCH 20/29] spec: remove --- CONTEXT.md | 53 -------- .../.openspec.yaml | 2 - .../design.md | 125 ------------------ .../proposal.md | 33 ----- .../manipulation-module-group-api/spec.md | 56 -------- .../specs/manipulation-operator/spec.md | 73 ---------- .../spec.md | 64 --------- .../spec.md | 45 ------- .../tasks.md | 48 ------- .../control-docs-housekeeping/.openspec.yaml | 2 - .../control-docs-housekeeping/README.md | 3 - .../control-docs-housekeeping/design.md | 26 ---- .../control-docs-housekeeping/proposal.md | 25 ---- .../specs/control-docs-housekeeping/spec.md | 15 --- .../control-docs-housekeeping/tasks.md | 17 --- .../viser-group-planning-ui/.openspec.yaml | 2 - .../changes/viser-group-planning-ui/README.md | 3 - .../changes/viser-group-planning-ui/design.md | 40 ------ .../viser-group-planning-ui/proposal.md | 28 ---- .../specs/viser-group-planning-ui/spec.md | 93 ------------- .../changes/viser-group-planning-ui/tasks.md | 17 --- openspec/config.yaml | 20 --- openspec/specs/group-aware-ik-rrt/spec.md | 38 ------ .../manipulation-module-group-api/spec.md | 27 ---- .../manipulation-planning-config/spec.md | 10 -- .../specs/planning-group-foundation/spec.md | 24 ---- .../planning-world-monitor-groups/spec.md | 26 ---- 27 files changed, 915 deletions(-) delete mode 100644 CONTEXT.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/design.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/proposal.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md delete mode 100644 openspec/changes/clean-manipulation-visualization-boundary/tasks.md delete mode 100644 openspec/changes/control-docs-housekeeping/.openspec.yaml delete mode 100644 openspec/changes/control-docs-housekeeping/README.md delete mode 100644 openspec/changes/control-docs-housekeeping/design.md delete mode 100644 openspec/changes/control-docs-housekeeping/proposal.md delete mode 100644 openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md delete mode 100644 openspec/changes/control-docs-housekeeping/tasks.md delete mode 100644 openspec/changes/viser-group-planning-ui/.openspec.yaml delete mode 100644 openspec/changes/viser-group-planning-ui/README.md delete mode 100644 openspec/changes/viser-group-planning-ui/design.md delete mode 100644 openspec/changes/viser-group-planning-ui/proposal.md delete mode 100644 openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md delete mode 100644 openspec/changes/viser-group-planning-ui/tasks.md delete mode 100644 openspec/config.yaml delete mode 100644 openspec/specs/group-aware-ik-rrt/spec.md delete mode 100644 openspec/specs/manipulation-module-group-api/spec.md delete mode 100644 openspec/specs/manipulation-planning-config/spec.md delete mode 100644 openspec/specs/planning-group-foundation/spec.md delete mode 100644 openspec/specs/planning-world-monitor-groups/spec.md diff --git a/CONTEXT.md b/CONTEXT.md deleted file mode 100644 index 2f418448a6..0000000000 --- a/CONTEXT.md +++ /dev/null @@ -1,53 +0,0 @@ -# DimOS Robotics - -DimOS composes robot capabilities into plans that can be previewed and executed across one or more robots. - -## Language - -**Planning**: -The complete process of finding a geometric motion and parameterizing it in time for preview and execution. -_Avoid_: Using “planning” to mean geometric path finding alone - -**Path**: -A geometric motion through configuration space without assigned timing, represented by planner-produced waypoints. -_Avoid_: Trajectory - -**Waypoint**: -A discrete robot configuration emitted by a planner as part of a path and consumed by trajectory parameterization. -_Avoid_: Trajectory point, path point - -**Planned Joint**: -A joint explicitly included in a generated plan’s waypoints and synchronized trajectory. A generated plan does not command joints outside this set. -_Avoid_: Every joint belonging to an affected robot - -**Generated Plan**: -The completed outcome of planning, containing both its geometric waypoints and one synchronized parameterized motion. An incomplete plan may exist transiently during construction but is not externally observable. -_Avoid_: Path, planner result - -**Plan Freshness**: -The condition that current planned-joint positions still match a generated plan’s synchronized trajectory start within the accepted tolerance. -_Avoid_: UI snapshot freshness, telemetry freshness - -**Synchronized Trajectory**: -A parameterized motion on one shared relative clock for every planned joint. -_Avoid_: Per-robot trajectory collection - -**Trajectory Parameterization**: -The assignment of timing and motion derivatives to a geometric waypoint path, producing the motion used by preview and execution. -_Avoid_: Preview sampling, path projection - -**Manipulation Operator**: -The human-facing control boundary attached to a manipulation visualization session, through which a client inspects manipulation state and requests target evaluation, planning, preview, execution, cancellation, or reset. It is shared by visualization clients rather than belonging to a particular renderer. -_Avoid_: Visualization backend, renderer - -**Target Draft**: -A frontend-owned, uncommitted set of desired joint or pose targets. A target draft is not a generated plan and has no execution authority. -_Avoid_: Plan, planned path - -**Target Evaluation**: -An advisory assessment of a target draft’s feasibility and resulting joint or pose state. It is feedback for an operator and is not planning or execution authority. -_Avoid_: Generated plan, validated plan - -**Preview**: -A transient visual playback of a generated plan’s synchronized trajectory. A preview has no execution authority and disappears when playback finishes or is cancelled. -_Avoid_: Target ghost, generated plan diff --git a/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml b/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml deleted file mode 100644 index 68b717479d..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-07-11 diff --git a/openspec/changes/clean-manipulation-visualization-boundary/design.md b/openspec/changes/clean-manipulation-visualization-boundary/design.md deleted file mode 100644 index 7805a3e4a8..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/design.md +++ /dev/null @@ -1,125 +0,0 @@ -## Context - -Viser currently receives `ManipulationModule` and `WorldMonitor` directly. Its GUI and visualization adapter call a broad set of module and monitor APIs for topology, current state, FK, collision checks, IK, planning, preview, execution, cancellation, and plan freshness. The Viser layer also accepts local and global joint-name variants, composes complete robot states, projects generated paths for preview, and duplicates safety policy that belongs in manipulation. - -The visualization protocol reinforces this coupling: backends pull current state, receive `GeneratedPlan`, resolve planning groups at runtime, and separately show, animate, cancel, and hide preview ghosts. Execution independently turns the stored geometric path into trajectories, so preview and execution need not share one timing authority. - -The agreed domain language in `CONTEXT.md` distinguishes paths, synchronized trajectories, generated plans, target drafts, target evaluations, plan freshness, manipulation operators, and transient previews. This design applies those distinctions while preserving the existing visible Viser interaction contract. - -## Goals / Non-Goals - -**Goals:** - -- Materialize one synchronized trajectory immediately after geometric planning and store it with the path in `GeneratedPlan`. -- Make the stored trajectory the sole motion and timing authority for preview and execution. -- Let the existing `JointTrajectoryTask` execute a configured subset without commanding omitted joints. -- Replace Viser's direct module/monitor dependencies with one concrete, UI-neutral `ManipulationOperator`. -- Make visualization topology immutable after initialization and push current joint-state frames instead of allowing renderers to pull world state. -- Pass raw globally named trajectories to visualization backends and keep projection backend-local. -- Move execution freshness and authoritative action validation out of Viser. -- Remove obsolete adapters and representation conversions rather than retaining compatibility shims. - -**Non-Goals:** - -- Changing planner algorithms, IK algorithms, collision ownership, or world-backend kinematics. -- Providing transport-level simultaneous multi-robot state snapshots; planning keeps the current static-world assumption and existing state APIs. -- Supporting dynamic robot or planning-group topology after visualization initialization. -- Adding atomic multi-task coordinator dispatch, dynamic subset resource claims, or concurrent joint ownership semantics. -- Correcting the existing lack of coordinator task cancellation when manipulation execution is cancelled. -- Changing Viser controls, wording, layout, group selection behavior, target ghosts, or pose-gizmo interactions. -- Adding arbitrary pose-target frame transforms; operator pose requests initially require the world frame. - -## Decisions - -### 1. Materialize `GeneratedPlan` once after planner success - -`PlannerSpec.plan` continues to produce the planner's geometric `PlanningResult`. `ManipulationModule` validates the successful path, resolves selected global joint limits by exact name and requested order, calls one trajectory generator across all planned joints, validates the result, and constructs a complete `GeneratedPlan` containing both `path` and `trajectory`. - -Only complete generated plans are cached or exposed. Preview, execution, status, and completion timing MUST NOT regenerate or re-parameterize the path. The path remains geometric/diagnostic data; the trajectory is motion/timing authority. - -This keeps planner backends waypoint-focused and avoids a transient incomplete `GeneratedPlan` contract. - -### 2. Store only planned global joints and no separate baseline - -Both `GeneratedPlan.path` and `GeneratedPlan.trajectory` use canonical global joint names in selected planning-group order. `path[0]` and the first trajectory point represent the planned-joint start state. Joints outside the selected groups are not stored or commanded. - -No plan baseline is added. Under the accepted single-commander/static-state assumption, renderers overlay planned joints onto their latest current visual state. Execution sends only planned joints through the existing trajectory task. - -### 3. Extend the existing `JointTrajectoryTask` for partial trajectories - -The task validates that incoming trajectory names are a non-empty, unique subset of its configured joints and that every point matches the active width with finite values and valid timing. During execution it emits `JointCommandOutput` using the incoming active names, not the task's full configured list. - -`claim()` continues to claim the full configured joint set. This intentionally preserves current arbitration/preemption behavior while preventing commands to omitted joints. Dynamic claims are deferred. - -### 4. Initialize a static visualization session - -`PlanningSceneInfo` is extended with resolved `PlanningGroup` values in addition to robot configs keyed by world robot ID. `ManipulationModule` constructs one concrete `ManipulationOperator` after world setup and asks `WorldMonitor` to initialize the visualization with a session containing: - -- immutable scene/robot/planning-group metadata; and -- the optional concrete operator used by interactive backends. - -There is no `ManipulationOperatorSpec` or protocol because only one implementation exists. `VisualizationSpec` remains justified by its multiple backend implementations. - -### 5. Push fast state and keep slow operator status small - -`WorldMonitor` periodically builds a `VisualizationStateFrame` containing current joint states keyed by world robot ID and calls `VisualizationSpec.update_state(frame)`. The frame contains no staleness policy. Viser uses it for current URDF rendering, target initialization, current presets, and joint readout. - -`ManipulationOperator.status()` returns only slow-changing manipulation state, error, and stored-plan summary. It does not return topology or current joint telemetry. - -### 6. Use a raw trajectory visualization API - -The runtime visualization API is reduced to raw globally named trajectory animation and cancellation. Backends receive no `GeneratedPlan`, path, or runtime group IDs. They derive affected robots from trajectory joint names using static initialization metadata, preserve stored timestamps by default, and apply an explicit duration only as playback-rate scaling without mutating the trajectory. - -Preview ghosts are transient. Animation shows affected ghosts, plays the trajectory, and hides them in `finally`; cancellation invalidates playback and hides ghosts immediately. `show_preview` and `hide_preview` are removed. - -### 7. Use one concrete stateless manipulation operator facade - -The Viser panel receives only `ManipulationOperator` and its local scene/controller objects. The facade composes `ManipulationModule` and `WorldMonitor` and exposes synchronous typed operations: - -- compact `status`; -- joint/pose target evaluation; -- joint/pose planning; -- preview, execute, cancel, clear-plan, and reset. - -Target drafts, group selection, and callback generation IDs remain frontend-owned. Operator evaluation and planning calls are stateless with respect to drafts. - -Joint requests contain ordered group IDs and one exact globally named `JointState` whose names equal the selected groups' concatenated joints in group order. Pose requests contain explicit world-frame `PoseStamped` targets, optional auxiliary group IDs, and an optional globally named seed. - -Target evaluations are advisory and expose only selected global joints, group poses, and group diagnostics. Complete per-robot states used for FK/collision checks stay private. Planning consumes the original target request rather than a prior evaluation result. Planning returns a typed summary, never the complete `GeneratedPlan`. - -### 8. Keep authoritative freshness in manipulation - -Viser no longer polls `is_state_stale`, records complete robot snapshots, compares snapshots, or decides whether execution is fresh. Immediately before dispatch, `ManipulationModule.execute` resolves current planned joints and compares them against the first stored trajectory point within the configured tolerance. Missing, malformed, or mismatched state rejects execution. - -Omitted joints are intentionally ignored under the accepted single-commander assumption. - -### 9. Route preview cancellation through the existing outbound port - -`ManipulationOperator.cancel` calls `ManipulationModule.cancel`. The module routes preview cancellation through `WorldMonitor`, which calls `VisualizationSpec.cancel_preview_animation`; no layer knows whether Viser or another backend implements the port. - -The actual coordinator execution-cancellation gap remains explicitly out of scope. - -## Risks / Trade-offs - -- **Partial trajectories change task validation/output behavior** → Preserve full task claims, add strict subset/shape/timing validation, and retain full-width trajectory tests. -- **No stored baseline means omitted joints can differ between preview moments** → Accepted under the stated static/single-commander assumption; planned joints and timing remain identical. -- **Raw global trajectories make renderers aware of global joint naming** → Supply immutable robot/group metadata once and centralize backend-local name projection helpers. -- **A visualization callback can re-enter the same backend through operator preview** → Keep operator calls synchronous but run them from Viser's existing worker lane; retain visualization mutation/cancellation locking and generation tests. -- **Removing frontend freshness checks changes immediate button behavior** → Keep authoritative module rejection and return clear action errors; do not duplicate policy in state frames. -- **Direct cutover touches several layers simultaneously** → Land model/materialization and task subset support first, then protocol/session/operator migration, then delete legacy methods and adapters after all backends compile. -- **Exact API names may need adjustment against construction code** → Preserve the directional contracts and responsibilities even if final concrete class or method names change during implementation. - -## Migration Plan - -1. Add trajectory materialization and partial-task support with focused tests while preserving existing visualization calls temporarily inside the implementation branch. -2. Add static visualization session/state-frame DTOs and the concrete operator; initialize them after world setup. -3. Migrate WorldMonitor, Drake, and Viser to `initialize`, `update_state`, raw `animate_trajectory`, and cancellation. -4. Migrate the Viser panel to typed operator requests/results and pushed current state; remove direct module/monitor access, snapshot freshness logic, and `GroupPanelBackend`. -5. Remove old visualization methods and all lazy path-to-trajectory/preview projection paths. -6. Run full manipulation, planning, visualization, control, lifecycle, and multi-robot tests plus mypy, Ruff, diff, and strict OpenSpec validation. - -Rollback is a source revert because this is an internal breaking boundary with no persisted data migration. - -## Open Questions - -- Final concrete DTO and method names are provisional until implementation verifies construction and lifecycle call sites; their responsibilities and directionality are fixed by this design. diff --git a/openspec/changes/clean-manipulation-visualization-boundary/proposal.md b/openspec/changes/clean-manipulation-visualization-boundary/proposal.md deleted file mode 100644 index 847ca42ac8..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/proposal.md +++ /dev/null @@ -1,33 +0,0 @@ -## Why - -The manipulation visualization boundary currently lets Viser pull from both `ManipulationModule` and `WorldMonitor`, so UI code performs planning-group lookup, state normalization, FK/collision composition, preview projection, and execution freshness checks. Preview and execution also derive different representations from an untimed stored path, making the boundary difficult to reason about and easy to drift. - -## What Changes - -- **BREAKING** Extend `GeneratedPlan` to contain both its geometric path and one synchronized, globally named `JointTrajectory`; parameterize the path exactly once immediately after `PlannerSpec.plan` succeeds and before the plan is cached or exposed. -- Make the stored trajectory the sole motion/timing authority for preview and execution. Preserve the geometric path for inspection and diagnostics. -- **BREAKING** Replace the pull-oriented visualization API with one-shot static initialization, pushed current-state frames, raw synchronized-trajectory animation, and cancellation. Remove `show_preview`, `hide_preview`, and `animate_plan(GeneratedPlan)`. -- Include resolved planning groups in immutable visualization initialization metadata. Visualization topology does not change during a session. -- Add one concrete `ManipulationOperator` facade that composes `ManipulationModule` and `WorldMonitor` for typed target evaluation and manipulation actions. Viser receives this facade through visualization initialization instead of receiving the raw module and monitor. -- Keep target drafts and asynchronous request sequencing frontend-owned. Make target evaluation advisory and require planning to consume the original complete target request. -- Move plan freshness enforcement into `ManipulationModule.execute`, comparing current planned joints with the stored trajectory start. Remove Viser-side robot snapshots, stale-state policy, and execution safety decisions. -- Push only current joint states to visualizations at visualization rate; freshness remains an authoritative manipulation concern and does not cross the rendering API. -- Update the existing `JointTrajectoryTask` to accept a non-empty configured joint subset while retaining its full configured resource claim. It emits commands only for trajectory joints, so omitted joints are not commanded. -- Preserve the existing Viser controls and interaction contract while making preview ghosts transient: completion or cancellation hides them immediately. - -## Capabilities - -### New Capabilities -- `manipulation-visualization-boundary`: Static visualization sessions, pushed joint-state frames, raw synchronized-trajectory preview, cancellation, and backend isolation from manipulation internals. -- `manipulation-operator`: The concrete UI-neutral facade, canonical target requests/evaluations, compact status, and action semantics used by interactive visualization panels. -- `partial-joint-trajectory-execution`: Validation and execution of configured joint subsets through the existing `JointTrajectoryTask` while preserving full resource claims. - -### Modified Capabilities -- `manipulation-module-group-api`: Generated plans become materialized path-plus-trajectory artifacts; preview and execution consume the stored trajectory and module-owned execution freshness. - -## Impact - -- Affects manipulation planning models and module materialization/execution, visualization protocols and scene metadata, `WorldMonitor`, Drake and Viser backends, Viser panel/backend/state code, trajectory control task validation/output, tests, and manipulation documentation. -- Removes direct `ManipulationModule` and `WorldMonitor` dependencies from Viser GUI/renderer construction. -- Replaces existing visualization method signatures and `GeneratedPlan` construction sites; this is an intentional internal breaking migration without compatibility shims. -- Does not change planner algorithms, robot/world collision ownership, coordinator batching, actual hardware execution cancellation, or the visible Viser interaction design. diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md deleted file mode 100644 index 58453778e5..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-module-group-api/spec.md +++ /dev/null @@ -1,56 +0,0 @@ -## ADDED Requirements - -### Requirement: Successful planning must materialize one complete generated plan -After `PlannerSpec.plan` returns a successful geometric path, ManipulationModule MUST validate the selected global waypoints and parameterize them exactly once into one synchronized globally named `JointTrajectory`. It MUST cache or expose a `GeneratedPlan` only when both path and trajectory are complete and mutually consistent. - -#### Scenario: Selected-group planning succeeds -- **WHEN** a planner returns valid waypoints for the requested planning groups -- **THEN** ManipulationModule stores one generated plan containing the original path and one trajectory over the same selected global joints - -#### Scenario: Parameterization fails -- **WHEN** joint limits, waypoint data, or generated trajectory validation fails -- **THEN** planning fails and no generated plan is cached or exposed - -#### Scenario: Multiple robots are selected -- **WHEN** selected group joints span multiple robots -- **THEN** one generator call creates one shared relative trajectory clock across all planned joints - -### Requirement: Stored trajectory must be preview and execution authority -Preview, execution, completion timing, and status logic MUST consume the stored generated-plan trajectory without re-projecting or re-parameterizing the geometric path. - -#### Scenario: Plan is previewed and executed -- **WHEN** a stored generated plan is previewed and later executed -- **THEN** both operations consume the same trajectory points, joint order, and timestamps - -#### Scenario: Preview duration is overridden -- **WHEN** a caller requests a different display duration -- **THEN** only visualization playback rate changes and the stored trajectory remains unchanged - -### Requirement: Generated trajectories must contain only planned joints -The generated path and trajectory MUST contain exactly the selected planning groups' global joint names in selection order. Joints outside the selected groups MUST NOT be added to the generated plan or execution command. - -#### Scenario: Planning group omits a robot joint -- **WHEN** a selected group excludes another configured robot joint -- **THEN** the excluded joint appears in neither the generated trajectory nor the trajectory task command - -### Requirement: Manipulation module must own execution freshness -Immediately before dispatch, ManipulationModule MUST compare current planned-joint positions with the stored trajectory's first point within configured tolerance. Missing, malformed, duplicate, reordered, stale, or mismatched planned-joint state MUST reject execution before task dispatch. - -#### Scenario: Planned joints still match -- **WHEN** all current planned joints match the trajectory start within tolerance -- **THEN** execution may translate and dispatch the stored trajectory - -#### Scenario: Planned joint moved after planning -- **WHEN** any current planned joint differs from the trajectory start beyond tolerance -- **THEN** execution is rejected without relying on frontend snapshots - -### Requirement: Execution must translate only at the coordinator boundary -Execution MUST split the stored global trajectory by affected robot without changing timing, translate selected local names to coordinator names at dispatch, and submit the resulting partial trajectory to the existing configured trajectory task. - -#### Scenario: Selected joints use coordinator aliases -- **WHEN** a robot config maps selected local joint names to coordinator names -- **THEN** dispatch preserves stored points and timestamps while translating only those selected names - -#### Scenario: Plan spans multiple robots -- **WHEN** planned global joints span multiple configured trajectory tasks -- **THEN** each task receives its selected local subset on the same stored relative clock using the existing sequential submission behavior diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md deleted file mode 100644 index 5ea9e4f106..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-operator/spec.md +++ /dev/null @@ -1,73 +0,0 @@ -## ADDED Requirements - -### Requirement: Interactive manipulation must use one concrete operator facade -The system MUST provide one concrete UI-neutral manipulation operator that composes `ManipulationModule` and `WorldMonitor`. It MUST NOT introduce an operator Spec or Protocol when no alternate implementation exists. - -#### Scenario: Visualization session is created -- **WHEN** ManipulationModule completes world initialization -- **THEN** it constructs one operator and binds that concrete instance into the visualization session - -### Requirement: Operator dynamic reads must remain compact -The operator MUST expose slow-changing manipulation state, error, and stored-plan summary through one compact status read. Static topology and high-rate joint states MUST NOT be returned by operator status. - -#### Scenario: Panel refreshes status -- **WHEN** Viser refreshes module and plan status -- **THEN** it reads compact operator status without copying robot configs or joint telemetry - -### Requirement: Joint target requests must use canonical global joints -Joint target requests MUST contain ordered planning-group IDs and one globally named `JointState` whose names exactly equal the selected groups' concatenated joint names in group order. Local aliases, suffix matching, missing joints, duplicate joints, and extraneous joints MUST be rejected. - -#### Scenario: Multiple groups submit a joint draft -- **WHEN** a frontend submits targets for ordered groups across one or more robots -- **THEN** the operator validates one flat global joint state in that exact selection order - -### Requirement: Pose target requests must carry explicit frames -Pose target requests MUST contain `PoseStamped` values keyed by pose-capable planning-group ID, optional ordered auxiliary group IDs, and an optional globally named seed. Initial operator support MUST require every pose target to use the world frame. - -#### Scenario: World-frame pose target is evaluated -- **WHEN** a frontend submits a world-frame target for a pose-capable group -- **THEN** the operator evaluates it without inventing or rewriting its frame - -#### Scenario: Unsupported pose frame is submitted -- **WHEN** a pose target uses a frame other than world -- **THEN** the operator rejects the request with a clear validation result - -### Requirement: Target evaluations must be advisory and selected-domain only -Target evaluation MUST return feasibility, status, message, selected globally named joints when available, group poses, and per-group diagnostics. Complete per-robot states used for FK or collision checks MUST remain internal, and an evaluation MUST NOT authorize planning or execution. - -#### Scenario: Joint target is evaluated -- **WHEN** the operator composes complete robot states for collision and FK checks -- **THEN** the frontend receives only selected-domain results and diagnostics - -#### Scenario: Pose target is evaluated -- **WHEN** IK produces selected joints -- **THEN** the operator returns those selected global joints and group results without exposing internal complete states - -### Requirement: Planning must consume original complete target requests -Operator planning methods MUST consume the original joint or pose target request and MUST NOT accept a prior advisory evaluation as planning authority. Planning MUST return a typed action summary rather than the complete `GeneratedPlan`. - -#### Scenario: Target changes after evaluation -- **WHEN** a frontend plans a newer target draft after an older evaluation completed -- **THEN** planning validates and plans the submitted request independently of the older evaluation - -#### Scenario: Planning succeeds -- **WHEN** manipulation caches a complete generated plan -- **THEN** the operator returns success, message, group IDs, waypoint count, and duration without returning path or trajectory data - -### Requirement: Target drafts and scheduling must remain frontend-owned -The operator MUST be synchronous and stateless with respect to editable target drafts, group selection, request generations, and frontend worker scheduling. - -#### Scenario: Evaluation results arrive out of order -- **WHEN** Viser runs synchronous operator calls in workers and a newer target generation exists -- **THEN** Viser discards the stale result using its own generation identity - -### Requirement: Operator actions must preserve domain ownership -Preview, execute, cancel, clear-plan, and reset methods MUST delegate to manipulation behavior and return typed action results. Cancellation MUST route preview cancellation through WorldMonitor's configured visualization without directly addressing Viser. - -#### Scenario: Operator cancels during preview -- **WHEN** the frontend calls operator cancel -- **THEN** ManipulationModule cancellation routes through WorldMonitor to visualization cancellation - -#### Scenario: Operator clears a plan -- **WHEN** the frontend calls clear-plan -- **THEN** the stored generated plan is discarded without introducing persistent preview state diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md deleted file mode 100644 index 0f7b46b791..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/specs/manipulation-visualization-boundary/spec.md +++ /dev/null @@ -1,64 +0,0 @@ -## ADDED Requirements - -### Requirement: Visualization initialization must receive immutable topology -The visualization protocol MUST initialize each backend once with robot configurations, resolved planning groups, and an optional concrete manipulation operator. Robot and planning-group topology MUST remain unchanged for the initialized session. - -#### Scenario: Interactive visualization initializes -- **WHEN** manipulation world setup completes with an interactive visualization backend -- **THEN** the backend receives all robot IDs, robot configs, resolved planning groups, and the concrete manipulation operator in one initialization session - -#### Scenario: Non-interactive visualization initializes -- **WHEN** a rendering backend does not provide an operator panel -- **THEN** it initializes from the same static topology and ignores the optional operator - -### Requirement: Visualization state must be pushed without world access -WorldMonitor MUST push current joint-state frames to visualization backends at visualization update cadence. A visualization backend MUST NOT pull current state, freshness, FK, collision, planning, or execution data from `WorldMonitor` or `ManipulationModule`. - -#### Scenario: Current robot state changes -- **WHEN** WorldMonitor publishes a visualization update -- **THEN** the backend receives current joint states keyed by initialized world robot IDs - -#### Scenario: State freshness changes -- **WHEN** telemetry becomes stale -- **THEN** no freshness policy is added to the visualization frame and authoritative manipulation actions validate state independently - -### Requirement: Visualization preview must consume raw synchronized trajectories -The visualization protocol MUST animate the stored globally named `JointTrajectory` directly and MUST NOT receive `GeneratedPlan` or regenerate timing from a geometric path. - -#### Scenario: Stored trajectory is previewed -- **WHEN** manipulation previews a generated plan without a display-duration override -- **THEN** the backend plays the raw trajectory using its stored timestamps - -#### Scenario: Display duration is overridden -- **WHEN** preview specifies a display duration -- **THEN** the backend scales playback delays without mutating or re-parameterizing the stored trajectory - -#### Scenario: Multiple robots are planned -- **WHEN** trajectory joint names span multiple initialized robots -- **THEN** the backend projects those joints through static topology and advances every affected preview ghost on the same stored clock - -### Requirement: Visualization preview must be transient -The visualization protocol MUST expose only trajectory animation and cancellation for preview lifecycle. Preview ghosts MUST disappear after completion, cancellation, error, replacement, or backend close. - -#### Scenario: Preview completes -- **WHEN** the final trajectory time is displayed -- **THEN** all preview ghosts are hidden immediately - -#### Scenario: Preview is cancelled -- **WHEN** cancellation occurs during playback -- **THEN** stale frames stop mutating the scene and all preview ghosts are hidden immediately - -#### Scenario: Preview is replaced -- **WHEN** a new trajectory begins while another preview is active -- **THEN** the old generation is cancelled and cannot mutate the new preview - -### Requirement: Visualization backends must remain domain-isolated -Viser GUI and renderer construction MUST NOT receive raw `ManipulationModule` or `WorldMonitor` dependencies. Runtime render methods MUST NOT perform planning-result validation, target evaluation, execution freshness checks, or action dispatch outside the concrete manipulation operator. - -#### Scenario: Viser panel starts -- **WHEN** the visualization session includes an operator -- **THEN** the panel uses only that operator for domain reads and actions - -#### Scenario: Preview renders -- **WHEN** Viser receives a valid raw trajectory -- **THEN** it performs only topology/name projection, visual-state overlay, playback, and ghost cleanup diff --git a/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md b/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md deleted file mode 100644 index 4f27ffa242..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/specs/partial-joint-trajectory-execution/spec.md +++ /dev/null @@ -1,45 +0,0 @@ -## ADDED Requirements - -### Requirement: Existing trajectory tasks must accept configured joint subsets -`JointTrajectoryTask` MUST accept a trajectory only when its joint names are a non-empty, unique subset of the task's configured joints and every trajectory point matches that active joint width. Unknown, duplicate, empty, malformed, non-finite, or invalidly timed input MUST be rejected before execution state changes. - -#### Scenario: Selected planning group trajectory is submitted -- **WHEN** a task configured for a complete arm receives a valid trajectory containing only that arm's selected planning-group joints -- **THEN** the existing task accepts and executes the trajectory - -#### Scenario: Unknown joint is submitted -- **WHEN** a trajectory contains a joint outside the task configuration -- **THEN** the task rejects it without entering executing state - -#### Scenario: Point width is malformed -- **WHEN** a trajectory point's positions or velocities do not match the trajectory joint-name count -- **THEN** the task rejects it before storing the trajectory - -### Requirement: Partial trajectory output must command only active joints -While executing a partial trajectory, `JointTrajectoryTask` MUST emit `JointCommandOutput` names and values only for the trajectory's active joint subset. It MUST NOT fill or command configured joints omitted from the trajectory. - -#### Scenario: Partial trajectory is sampled -- **WHEN** a two-joint trajectory executes on a six-joint configured task -- **THEN** each output contains exactly those two joint names and two sampled positions - -#### Scenario: Full trajectory is sampled -- **WHEN** a trajectory contains every configured task joint -- **THEN** existing full-width output behavior remains unchanged - -### Requirement: Resource claims must remain configuration-wide -`JointTrajectoryTask.claim()` MUST continue to claim the task's full configured joint set regardless of the active trajectory subset. - -#### Scenario: Partial trajectory executes -- **WHEN** the active trajectory omits configured joints -- **THEN** arbitration and preemption still use the task's full configured resource claim - -### Requirement: Partial trajectory lifecycle must clear active names safely -Replacement, completion, cancellation, reset, and fault paths MUST keep active trajectory joint names consistent with the stored trajectory and MUST NOT leak names from a prior execution. - -#### Scenario: Partial trajectory is replaced -- **WHEN** an active subset trajectory is replaced by a different valid subset -- **THEN** subsequent outputs use only the replacement trajectory's names - -#### Scenario: Task resets -- **WHEN** a completed or aborted task resets -- **THEN** its stored trajectory and active subset are cleared while its configured claim remains unchanged diff --git a/openspec/changes/clean-manipulation-visualization-boundary/tasks.md b/openspec/changes/clean-manipulation-visualization-boundary/tasks.md deleted file mode 100644 index 323099d4e8..0000000000 --- a/openspec/changes/clean-manipulation-visualization-boundary/tasks.md +++ /dev/null @@ -1,48 +0,0 @@ -## 1. Materialize Generated Plans - -- [x] 1.1 Extend `GeneratedPlan` with a required synchronized `JointTrajectory`, migrate construction sites, and add model/import/round-trip tests proving path and trajectory remain distinct representations. -- [x] 1.2 Add one post-planner materialization path in `ManipulationModule` that validates selected global waypoints, resolves velocity/acceleration limits by exact global name and selection order, and invokes one trajectory generator across all planned joints. -- [x] 1.3 Validate generated trajectory names, dimensions, finite values, timestamps, positive duration, and ordered waypoint boundaries before caching; ensure planning/parameterization failure exposes no generated plan. -- [x] 1.4 Remove non-dispatch lazy path-to-trajectory generation from status/completion timing and isolate the remaining dispatch-only bridge for Phase 2 removal after partial trajectory task support lands. -- [x] 1.5 Add materialization tests for reordered groups, heterogeneous multi-robot limits, malformed/non-finite waypoints, invalid limits, generator failure, exactly one generation call, shared timing, and zero generation after caching. - -## 2. Execute Partial Joint Trajectories - -- [x] 2.1 Update the existing `JointTrajectoryTask.execute` validation to accept only non-empty unique configured subsets with valid point widths, finite values, and strictly increasing timing while preserving full-width behavior. -- [x] 2.2 Make task computation emit only the active trajectory joint names and sampled values while leaving `claim()` on the full configured joint set and clearing active subset state on replacement, completion, cancellation, fault, and reset. -- [x] 2.3 Update manipulation execution to split the stored global trajectory by robot, preserve every timestamp/value, translate selected names only at coordinator dispatch, and never fill or command omitted joints. -- [x] 2.4 Add control/manipulation tests for valid subsets, unknown/duplicate/empty names, malformed positions/velocities/timing, full-width compatibility, full configured claims, replacement/reset lifecycle, coordinator translation, omitted-joint non-commanding, and multi-robot shared clocks. - -## 3. Introduce the Visualization Session Boundary - -- [x] 3.1 Extend `PlanningSceneInfo` with resolved immutable planning groups and add visualization session/state-frame DTOs containing static topology, the optional concrete operator, and pushed current joint states respectively. -- [x] 3.2 Replace `VisualizationSpec` pull/group/plan methods with one-shot initialization, `update_state`, raw globally named `animate_trajectory`, cancellation, URL, and close; add structural/import tests for the complete protocol. -- [x] 3.3 Make `WorldMonitor` initialize visualization once, push current joint-state frames at visualization cadence, serialize renderer mutations, forward raw stored trajectories, and route cancellation without exposing freshness policy. -- [x] 3.4 Migrate Drake/Meshcat visualization to static topology and raw trajectory playback, preserving stored non-uniform timestamps, optional playback scaling, synchronized multi-robot ticks, generation cancellation, and automatic ghost cleanup. -- [x] 3.5 Add WorldMonitor and Drake tests for immutable group metadata, pushed state frames, no state pull, raw timing/scaling, trajectory immutability, multi-robot shared ticks, replacement/cancel/error/close cleanup, and removal of legacy visualization methods. - -## 4. Add the Concrete Manipulation Operator - -- [x] 4.1 Add concrete operator status, canonical joint/pose request, target evaluation, plan summary, action result, and visualization session types without introducing a new operator Protocol or Spec. -- [x] 4.2 Implement exact ordered global-joint request validation and explicit world-frame pose request validation, rejecting aliases, duplicates, omissions, extras, unsupported groups/frames, and malformed seeds. -- [x] 4.3 Move joint/pose target composition, IK, FK, per-robot state validity, and selected-domain evaluation results from Viser-local helpers into the operator while keeping complete robot states private. -- [x] 4.4 Implement synchronous stateless planning methods that consume original requests, delegate to manipulation APIs, and return summaries without exposing `GeneratedPlan`; implement compact slow-changing status and action methods. -- [x] 4.5 Move execute freshness enforcement into `ManipulationModule` by comparing current planned global joints with the stored trajectory start immediately before dispatch; remove dependence on frontend robot snapshots and stale-state gates. -- [x] 4.6 Construct one operator after module/world setup, bind it into visualization initialization, preserve indirect module-to-world-to-visualization cancellation, and add lifecycle/action/freshness tests. - -## 5. Migrate Viser to the Clean Boundary - -- [x] 5.1 Initialize Viser from static robot/group metadata and build deterministic global-joint-to-robot/local-joint maps without retaining `ManipulationModule` or `WorldMonitor` dependencies. -- [x] 5.2 Consume pushed state frames for current URDFs, joint readout, target initialization, and Current presets; remove direct state/freshness polling and Viser-side stale policy. -- [x] 5.3 Replace generated-plan preview projection and preview-track resampling with raw stored-trajectory projection/playback; support duration scaling only, hide ghosts on completion/cancel/error/replacement/close, and prevent stale frames from mutating scene handles. -- [x] 5.4 Migrate the panel to the concrete operator's typed status, target evaluation, planning, preview, execution, cancellation, clear, and reset calls while keeping target drafts and callback generations frontend-owned. -- [x] 5.5 Remove `GroupPanelBackend`, direct module/monitor calls, local/global alias normalization, complete-state composition, robot snapshot matching, and obsolete preview show/hide state from Viser code. -- [x] 5.6 Preserve the normative Viser UI controls and interactions exactly and add granular tests for group selection, presets, sliders, pose gizmos, advisory evaluation, action results, pushed current state, raw preview timing, auto-hide, cancellation races, and absence of forbidden dependencies. - -## 6. Cleanup, Documentation, and Validation - -- [x] 6.1 Remove all production/test references to `animate_plan`, `show_preview`, `hide_preview`, lazy path projection/parameterization, Viser `ManipulationModule`/`WorldMonitor` access, and legacy panel adapter APIs. -- [x] 6.2 Update manipulation/planning-group/visualization documentation and `CONTEXT.md` to describe materialized generated plans, partial planned-joint commands, the concrete operator, pushed visualization state, and transient previews. -- [x] 6.3 Run focused planner/materialization, manipulation module/unit/reservation, WorldMonitor, Drake, Viser/factory/lifecycle, control task/coordinator, and multi-robot end-to-end tests with no relevant skips. -- [x] 6.4 Run targeted mypy over every changed production file, Ruff check/format, `git diff --check`, and `openspec validate clean-manipulation-visualization-boundary --strict`; audit the final diff for unrelated planner/control/UI changes. -- [x] 6.5 Complete a final architecture and safety review confirming one trajectory-generation call, stored-timing preview/execution parity, partial-output/full-claim semantics, module-owned freshness, renderer isolation, exact Viser interaction parity, and no compatibility shims. diff --git a/openspec/changes/control-docs-housekeeping/.openspec.yaml b/openspec/changes/control-docs-housekeeping/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/control-docs-housekeeping/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/control-docs-housekeeping/README.md b/openspec/changes/control-docs-housekeeping/README.md deleted file mode 100644 index ee333ffc5a..0000000000 --- a/openspec/changes/control-docs-housekeeping/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# control-docs-housekeeping - -Extract PR 6 from reference branch cc/spec/movegroup: peel off control/coordinator/generated/docs/housekeeping changes not central to planning-group semantics. diff --git a/openspec/changes/control-docs-housekeeping/design.md b/openspec/changes/control-docs-housekeeping/design.md deleted file mode 100644 index c42efe8508..0000000000 --- a/openspec/changes/control-docs-housekeeping/design.md +++ /dev/null @@ -1,26 +0,0 @@ -## Context - -These files do not obviously belong to the core planning-group review. Some may be genuine integration requirements; others may be incidental edits from the large branch. - -## Goals / Non-Goals - -**Goals:** -- Separate orthogonal changes from the planning-group PR stack. -- Keep generated files and docs synchronized only when required. -- Provide a clear report of accidental or questionable files. - -**Non-Goals:** -- Do not change planning-group semantics. -- Do not include Viser UI or planner algorithm changes. -- Do not include files just because they exist on the reference branch. - -## Decisions - -- Treat this change as a triage/extraction PR, not a mandatory feature PR. -- Include control/task changes only if validation shows they are required or independently useful. -- Keep generated registry changes next to the PR that introduces the corresponding blueprint, unless generation must happen at the end of the stack. - -## Risks / Trade-offs - -- If control changes depend on module API changes, this PR should stack after PR 4. -- If docs describe earlier PR behavior, those docs should move into the earlier PRs instead of this housekeeping PR. diff --git a/openspec/changes/control-docs-housekeeping/proposal.md b/openspec/changes/control-docs-housekeeping/proposal.md deleted file mode 100644 index 436f94fb41..0000000000 --- a/openspec/changes/control-docs-housekeeping/proposal.md +++ /dev/null @@ -1,25 +0,0 @@ -## Why - -The current branch contains small control, generated-registry, docs, and housekeeping changes mixed into the planning-group refactor. These should be peeled off so reviewers can decide whether they are required or accidental. - -## What Changes - -- Extract control/coordinator task and tick-loop changes if they are truly required. -- Extract generated blueprint registry updates if they are caused by earlier PRs. -- Extract remaining docs/readme cleanup and packaging/config touches. -- Report any file that appears accidental rather than including it blindly. - -## Capabilities - -### New Capabilities -- `control-docs-housekeeping`: Distribution plan for orthogonal control, generated, docs, and repo housekeeping changes. - -### Modified Capabilities -- `control-coordinator-integration`: Control/coordinator behavior updates only if validated as required by the planning-group stack. - -## Impact - -- Base branch: preferably PR 4 `manipulation-module-group-api`, or `main` for truly independent control/docs changes. -- Reference implementation: `cc/spec/movegroup`. -- Primary candidate files: `dimos/control/*`, `dimos/e2e_tests/test_control_coordinator.py`, generated blueprint files, docs/readmes, `pyproject.toml`, `AGENTS.md`. -- Out of scope: planning-group core, planners, module API, and Viser UI. diff --git a/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md b/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md deleted file mode 100644 index cc200026e1..0000000000 --- a/openspec/changes/control-docs-housekeeping/specs/control-docs-housekeeping/spec.md +++ /dev/null @@ -1,15 +0,0 @@ -## ADDED Requirements - -### Requirement: Orthogonal changes must be reviewed separately -Changes that are not necessary for planning-group semantics MUST be extracted into a separate review or explicitly excluded from the stack. - -#### Scenario: Candidate file is unrelated to planning groups -- **WHEN** a file change is docs-only, generated-only, control-only, or repo-housekeeping-only -- **THEN** it is reviewed in this housekeeping change or reported as intentionally excluded - -### Requirement: Generated files must match their source changes -Generated registry files MUST be included only with the source blueprint/config change that requires regeneration, or in a final generated-files PR. - -#### Scenario: Blueprint registry changes are generated -- **WHEN** a generated registry file differs from `main` -- **THEN** the extraction task identifies the source change and validates regeneration diff --git a/openspec/changes/control-docs-housekeeping/tasks.md b/openspec/changes/control-docs-housekeeping/tasks.md deleted file mode 100644 index 6831183321..0000000000 --- a/openspec/changes/control-docs-housekeeping/tasks.md +++ /dev/null @@ -1,17 +0,0 @@ -## 1. Triage candidate files - -- [ ] 1.1 Compare each candidate file against the reference branch and decide whether it is required, independent, generated, docs-only, or accidental. -- [ ] 1.2 Extract only required or intentionally independent files. -- [ ] 1.3 Move docs back into earlier PRs when they explain earlier PR behavior. - -## 2. Control/generated/docs extraction - -- [ ] 2.1 Extract control/task changes if they are intentional. -- [ ] 2.2 Extract generated blueprint registry changes only when their source blueprint is included. -- [ ] 2.3 Extract `pyproject.toml` and `AGENTS.md` only if justified; otherwise report as excluded. - -## 3. Validation - -- [ ] 3.1 Run `uv run pytest dimos/control/test_control.py dimos/e2e_tests/test_control_coordinator.py dimos/robot/test_all_blueprints.py -q` for included files. -- [ ] 3.2 Run any docs or generated-registry checks required by included changes. -- [ ] 3.3 Return a file-by-file inclusion/exclusion summary. diff --git a/openspec/changes/viser-group-planning-ui/.openspec.yaml b/openspec/changes/viser-group-planning-ui/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/viser-group-planning-ui/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/viser-group-planning-ui/README.md b/openspec/changes/viser-group-planning-ui/README.md deleted file mode 100644 index 72628be56b..0000000000 --- a/openspec/changes/viser-group-planning-ui/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# viser-group-planning-ui - -Extract PR 5 from reference branch cc/spec/movegroup: make Viser visualization and panel workflows group-aware. diff --git a/openspec/changes/viser-group-planning-ui/design.md b/openspec/changes/viser-group-planning-ui/design.md deleted file mode 100644 index 089ebc6a17..0000000000 --- a/openspec/changes/viser-group-planning-ui/design.md +++ /dev/null @@ -1,40 +0,0 @@ -## Context - -Viser must let users select planning groups, move joint and pose targets, see feasibility feedback, preview paths, and execute only when the preview still matches the current robot state. - -## Goals / Non-Goals - -**Goals:** -- Expose group selection and group-aware robot controls. -- Keep target ghost and preview animation tied to selected groups. -- Preserve safe execution behavior and clear recoverable errors. -- Keep UI review isolated from planner correctness review. - -**Non-Goals:** -- Do not change core planning algorithms in this PR. -- Do not alter the planning-group data model. -- Do not include control/coordinator changes. - -## Decisions - -- Use a panel backend boundary so UI code does not directly own all manipulation-module orchestration. -- Treat planning-group IDs as the UI's stable selection keys. -- Validate base-link/root assumptions before rendering URDFs under base poses. -- Keep manual demo/checklist coverage for UI feel in addition to unit tests. -- Treat `cc/spec/movegroup@0edb8d3dd` as the sole source for visible controls and interaction behavior; do not derive behavior from the upstream extraction commit. -- Adapt the reference UI to PR 4's existing explicit group APIs rather than importing manipulation or planner changes from the reference branch. -- Keep execution trajectory parameterization in the manipulation/module layer. Group-native preview renderers may map canonical `GeneratedPlan` waypoints onto the requested display duration, but must not reinterpret them as execution timing. -- Use the reference group-native visualization contract: `show_preview(Sequence[PlanningGroupID])`, `hide_preview(Sequence[PlanningGroupID])`, and `animate_plan(GeneratedPlan, duration: float = 3.0)`. `ManipulationModule.preview_plan` supplies `1.0` when duration is omitted, matching the normative reference. Update Viser and the existing Meshcat implementation together; do not import unrelated source-branch planning changes. -- Treat preview as one validated transaction. Resolve and deduplicate affected robots, snapshot each robot's non-selected joints once, project selected global joints into complete local frames from that fixed baseline, reject malformed/unknown/incomplete input before showing ghosts, then update every ghost once per shared tick before sleeping once. Unequal renderer frame counts map onto the same normalized clock. -- Use an animation generation/cancellation token. Cancel, clear, close, or a replacement preview invalidates the old generation so it cannot mutate removed or replacement scene handles. Serialize preview-scene mutation against periodic visualization publication. -- Group selection owns a monotonically increasing epoch. Target/plan callbacks commit only if both their operation sequence and selection epoch still match. Execute snapshots every affected robot immediately before dispatch, compares every selected group's ordered joints, and calls `ManipulationModule.execute()` without a robot filter. -- Execute snapshots every affected robot immediately before dispatch and delegates the stored plan without a robot filter. An external caller can unavoidably replace that stored plan between the panel freshness check and dispatch. -- Do not copy absent reference APIs such as `ManipulationModule.forward_kinematics` or `check_collision`. Use PR4 `WorldMonitor.get_group_ee_pose` for group FK and `WorldMonitor.is_state_valid` on composed full per-robot target states. Multi-robot validity is a conjunction of per-robot checks, not proof of simultaneous target-target collision freedom; pose IK retains PR4's coordinated collision check. - -## Risks / Trade-offs - -- UI tests can pass while interaction feels wrong; require a small manual checklist before review. -- This PR is still large, but isolating it keeps visual review focused. -- A whole-commit cherry-pick would leak unrelated planner and module changes; extraction must be semantic and visualization-scoped. -- The preview contract is the sole approved scope expansion outside visualization. Keep its module/monitor/protocol/world changes minimal and independently testable. -- A compatibility `robot_name` argument may validate that the named robot is affected, but MUST NOT trim a multi-robot generated plan before preview. diff --git a/openspec/changes/viser-group-planning-ui/proposal.md b/openspec/changes/viser-group-planning-ui/proposal.md deleted file mode 100644 index 99079912d3..0000000000 --- a/openspec/changes/viser-group-planning-ui/proposal.md +++ /dev/null @@ -1,28 +0,0 @@ -## Why - -The Viser changes are large and have different review criteria from planning algorithms. They should be reviewed as a dedicated UI/backend PR after the public manipulation group APIs exist. - -## What Changes - -- Make the Viser panel group-aware. -- Split panel/backend behavior out of the old adapter shape. -- Update scene previews, target ghosts, group selection, feasibility state, and safe execution checks. -- Update Viser tests and manual review checklist. - -## Capabilities - -### New Capabilities -- `viser-group-planning-ui`: Viser visualization supports group-aware planning, preview, target evaluation, and execution state. - -### Modified Capabilities -- `manipulation-visualization`: Visualization targets explicit planning groups rather than a robot-scoped end-effector field. - -## Impact - -- Base branch: PR 4 `manipulation-module-group-api`. -- Normative UI reference: `cc/spec/movegroup@0edb8d3dd`. User-visible deviations require explicit approval. -- API base: PR 4 at `cc/planning_group/main@05c25787a`; its explicit planning-group API scope remains unchanged. -- The upstream extraction commit `origin/cc/planning_group/viser@737fb3381` is not an implementation source. -- Primary files: `dimos/manipulation/visualization/viser/*`, `dimos/manipulation/visualization/types.py`, `dimos/manipulation/visualization/test_factory.py`. -- Supporting preview seam: group-native visualization protocol, world-monitor forwarding, manipulation-module preview routing, and the existing Meshcat implementation. -- Out of scope: all other manipulation-module, planning model/backend/algorithm, and control changes except as already supplied by earlier PRs. diff --git a/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md b/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md deleted file mode 100644 index f28750806f..0000000000 --- a/openspec/changes/viser-group-planning-ui/specs/viser-group-planning-ui/spec.md +++ /dev/null @@ -1,93 +0,0 @@ -## ADDED Requirements - -### Requirement: Viser panel must select planning groups explicitly -The Viser planning panel MUST present and use planning-group selections for group-aware joint and pose planning controls. - -#### Scenario: User selects a manipulator group -- **WHEN** a user selects a planning group in the panel -- **THEN** joint controls, pose target controls, and preview state apply to that group - -### Requirement: Viser panel must preserve the normative reference interaction contract -The Viser panel MUST match the controls, labels, ordering, defaults, selection behavior, presets, joint sliders, pose gizmos, target ghosts, feasibility colors, preview behavior, execute gate, and clear behavior from `cc/spec/movegroup@0edb8d3dd`. - -#### Scenario: Extracted panel is reviewed -- **WHEN** the extracted panel is compared with the normative reference -- **THEN** no user-visible deviation is present unless separately approved - -### Requirement: Viser must preserve the cleaned planning-group API boundary -The Viser implementation MUST consume PR 4's explicit planning-group APIs without adding planner APIs. The only manipulation/visualization API change permitted is the group-native preview transaction required below. - -#### Scenario: Backend adaptation is required -- **WHEN** reference UI code expects a different backend shape -- **THEN** the visualization backend translates the call without changing planner algorithms or control production code - -### Requirement: Preview must be group-native and synchronized -The visualization protocol MUST expose `show_preview` and `hide_preview` over ordered planning-group IDs and `animate_plan` over one generated plan plus a total display duration. It MUST preview the plan as one validated transaction so all affected robots advance on a shared animation clock. - -#### Scenario: Selected groups span multiple robots -- **WHEN** the user previews a fresh multi-robot plan -- **THEN** all affected preview ghosts advance together from the same generated plan - -#### Scenario: Preview duration is omitted -- **WHEN** `ManipulationModule.preview_plan` is called without a duration -- **THEN** it sends the full generated plan to `animate_plan` with a total display duration of `1.0` second - -#### Scenario: Visualization is called directly -- **WHEN** `animate_plan` is called without an explicit duration -- **THEN** its total display duration defaults to `3.0` seconds - -#### Scenario: A plan controls a subset of robot joints -- **WHEN** a generated plan contains selected global joints -- **THEN** each affected robot's unselected joints come from one fixed pre-playback snapshot for every animation frame - -#### Scenario: Preview input is invalid -- **WHEN** a plan contains unknown groups, malformed joint states, or cannot form complete affected-robot frames -- **THEN** preview rejects the transaction before showing any ghost - -#### Scenario: Preview is replaced or stopped -- **WHEN** cancel, clear, close, or a new preview invalidates an active preview generation -- **THEN** old animation frames no longer mutate scene handles - -### Requirement: Panel freshness must cover the full selected target set -The panel MUST associate asynchronous results with a selection epoch and MUST snapshot ordered joints for every affected robot when a plan becomes executable. - -#### Scenario: Selection changes during evaluation -- **WHEN** a planning group is toggled while target evaluation is in flight -- **THEN** the stale result does not update target state or visuals - -#### Scenario: A newer operation supersedes an older operation -- **WHEN** callbacks arrive out of order for the same selection epoch -- **THEN** only the callback whose operation sequence is current may update panel or scene state - -#### Scenario: Execute spans multiple robots -- **WHEN** every selected group's ordered joints still match the stored all-robot snapshot -- **THEN** execute delegates the full stored plan without applying a robot filter - -#### Scenario: Another caller replaces the stored module plan -- **WHEN** replacement occurs between the panel freshness check and execution -- **THEN** the external-plan replacement race is unavoidable - -### Requirement: Joint-target feasibility must use existing PR 4 world APIs -The visualization backend MUST compose full per-robot target states and call `WorldMonitor.is_state_valid`; it MUST NOT add or copy manipulation-module FK or collision APIs. - -#### Scenario: Multiple robots have joint targets -- **WHEN** each composed per-robot state is valid -- **THEN** the panel reports the conjunction of those checks without claiming simultaneous target-target collision proof - -#### Scenario: Pose targets are evaluated -- **WHEN** selected pose groups and auxiliary groups form a target set -- **THEN** the backend uses PR 4 `inverse_kinematics` with coordinated collision checking, then validates each affected robot's composed full target state through `WorldMonitor.is_state_valid` - -### Requirement: Viser preview must reflect group feasibility and current state -The Viser UI MUST show target feasibility and preview validity for the selected planning group. - -#### Scenario: Target is infeasible -- **WHEN** target evaluation reports infeasible for the selected group -- **THEN** the target ghost and panel state indicate that the target cannot be executed safely - -### Requirement: Viser execution must require a fresh matching plan -The Viser panel MUST prevent execution when no fresh plan matches the current selected robot/group state. - -#### Scenario: Current state no longer matches preview -- **WHEN** the robot state changes after planning -- **THEN** the panel rejects execute until all selected-group snapshots match again diff --git a/openspec/changes/viser-group-planning-ui/tasks.md b/openspec/changes/viser-group-planning-ui/tasks.md deleted file mode 100644 index b9a75ae999..0000000000 --- a/openspec/changes/viser-group-planning-ui/tasks.md +++ /dev/null @@ -1,17 +0,0 @@ -## 1. Extract Viser UI changes - -- [x] 1.1 Start from PR 4 and use only `cc/spec/movegroup@0edb8d3dd` as the normative UI reference; do not rely on the upstream extraction commit. -- [x] 1.2 Extract Viser panel/backend/scene/state/visualizer group-aware changes and the minimal group-native shared-clock preview seam. -- [x] 1.3 Remove or replace the old adapter code only as required by this UI slice. -- [x] 1.4 Extract visualization type/factory updates needed by Viser. - -## 2. Tests and manual check - -- [x] 2.1 Bring over Viser unit/lifecycle tests and add protocol, monitor forwarding, module routing, Meshcat projection, shared-tick, all-robot freshness, and preview cancellation race coverage. -- [x] 2.2 Run the manual checklist: group selector, joint sliders, pose gizmo, target ghost, infeasible target color, path preview, execute gate, clear path. - -## 3. Validation - -- [x] 3.1 Run `uv run pytest dimos/manipulation/test_manipulation_unit.py dimos/manipulation/planning/monitor/test_world_monitor.py dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/visualization/test_factory.py dimos/manipulation/visualization/viser/test_*.py -q` and `openspec validate viser-group-planning-ui --strict`. -- [x] 3.2 Run mandatory targeted mypy on changed Viser production files plus `manipulation_module.py`, `protocols.py`, `world_monitor.py`, and `drake_world.py`. -- [x] 3.3 Verify changes outside visualization are limited to the approved group-native preview protocol/forwarding/implementations, with no planner/backend algorithm or control changes. diff --git a/openspec/config.yaml b/openspec/config.yaml deleted file mode 100644 index 392946c67c..0000000000 --- a/openspec/config.yaml +++ /dev/null @@ -1,20 +0,0 @@ -schema: spec-driven - -# Project context (optional) -# This is shown to AI when creating artifacts. -# Add your tech stack, conventions, style guides, domain knowledge, etc. -# Example: -# context: | -# Tech stack: TypeScript, React, Node.js -# We use conventional commits -# Domain: e-commerce platform - -# Per-artifact rules (optional) -# Add custom rules for specific artifacts. -# Example: -# rules: -# proposal: -# - Keep proposals under 500 words -# - Always include a "Non-goals" section -# tasks: -# - Break tasks into chunks of max 2 hours diff --git a/openspec/specs/group-aware-ik-rrt/spec.md b/openspec/specs/group-aware-ik-rrt/spec.md deleted file mode 100644 index e6ea4839f5..0000000000 --- a/openspec/specs/group-aware-ik-rrt/spec.md +++ /dev/null @@ -1,38 +0,0 @@ -## Purpose - -Define the planning-group-aware IK and selected-joint planning algorithm contracts. - -## Requirements - -### Requirement: IK solvers must use planning-group target frames -IK solvers MUST resolve pose target frames from the requested planning group's `tip_link` rather than from robot-scoped end-effector metadata. - -`solve_pose_targets` is a pose-target set IK API: callers MAY provide one or more pose targets keyed by planning group. Backends MUST either solve the requested target set or return `IKStatus.UNSUPPORTED` when the requested target-set shape is outside that backend's capability. - -#### Scenario: Pose IK targets a manipulator group -- **WHEN** a pose target is submitted for a group with `tip_link="tcp"` -- **THEN** the IK solver constrains the `tcp` frame - -#### Scenario: Pink IK receives multiple pose targets -- **WHEN** Pink IK receives multiple pose targets for pose-targetable planning groups -- **THEN** it solves the target frames as a target set and returns selected joints in planning-group selection order - -#### Scenario: Single-target IK backend receives multiple pose targets -- **WHEN** Drake optimization IK or Jacobian IK receives more than one pose target or auxiliary planning groups -- **THEN** it returns `IKStatus.UNSUPPORTED` with an explanatory message - -### Requirement: Planners must preserve group-local joint ordering -Planners MUST accept and return joint targets in the requested group's local joint order while projecting through full robot state for collision checks. - -Selected joint planning MUST be part of the planner protocol. Planner backends that do not support selected/group planning MUST return `PlanningStatus.UNSUPPORTED` with an explanatory message rather than treating unsupported capability as a failed search. - -#### Scenario: Group joint target uses subset order -- **WHEN** a group target names a subset of robot joints -- **THEN** planning uses the correct full robot state and returns a path scoped to the group - -### Requirement: Algorithms must fail clearly for non-pose-targetable groups -Algorithms that require a pose target frame MUST return an explicit unsupported failure when the requested group has no `tip_link`. - -#### Scenario: Pose IK targets a joint-only group -- **WHEN** a pose target is submitted for a group without `tip_link` -- **THEN** the solver returns `IKStatus.UNSUPPORTED` with an explanatory message diff --git a/openspec/specs/manipulation-module-group-api/spec.md b/openspec/specs/manipulation-module-group-api/spec.md deleted file mode 100644 index b812df591f..0000000000 --- a/openspec/specs/manipulation-module-group-api/spec.md +++ /dev/null @@ -1,27 +0,0 @@ -## Purpose - -Define the public manipulation module APIs for explicit planning-group targeting and safe robot-scoped compatibility behavior. - -## Requirements - -### Requirement: Manipulation module must expose explicit group planning APIs -The manipulation module MUST allow callers to target planning groups explicitly for joint and pose planning operations. - -#### Scenario: Caller plans to a group pose -- **WHEN** a caller submits a pose target for a valid group ID -- **THEN** the module plans using that group and returns the planning result through the existing public contract - -### Requirement: Robot-scoped pose wrappers must fail safely without a unique pose group -Robot-scoped pose wrapper APIs MUST fail safely when a selected robot has zero or multiple pose-targetable groups. - -#### Scenario: End-effector pose is requested with no unique pose group -- **WHEN** `get_ee_pose` cannot resolve exactly one pose-targetable group -- **THEN** it returns `None` - -#### Scenario: Robot-scoped pose planning is requested with no unique pose group -- **WHEN** `plan_to_pose` cannot resolve exactly one pose-targetable group -- **THEN** it returns `False` - -#### Scenario: Robot-scoped IK is requested with no unique pose group -- **WHEN** `inverse_kinematics_single` cannot resolve exactly one pose-targetable group -- **THEN** it returns an `IKResult` with `NO_SOLUTION` diff --git a/openspec/specs/manipulation-planning-config/spec.md b/openspec/specs/manipulation-planning-config/spec.md deleted file mode 100644 index bd13e49196..0000000000 --- a/openspec/specs/manipulation-planning-config/spec.md +++ /dev/null @@ -1,10 +0,0 @@ -# manipulation-planning-config Specification - -## Requirements - -### Requirement: Robot model configuration must separate placement from planning chains -Robot model configuration MUST treat robot-scoped `base_link` as the link placed by `base_pose` and used by backend weld/strip behavior. Planning-chain base links and pose target frames MUST come from planning groups. - -#### Scenario: Robot has one pose-targetable group -- **WHEN** a robot model has `base_link="base_link"` and a planning group with `base_link="link1"` and `tip_link="tcp"` -- **THEN** placement uses the robot-scoped `base_link` while pose planning targets `tcp` diff --git a/openspec/specs/planning-group-foundation/spec.md b/openspec/specs/planning-group-foundation/spec.md deleted file mode 100644 index 6124a3bf5c..0000000000 --- a/openspec/specs/planning-group-foundation/spec.md +++ /dev/null @@ -1,24 +0,0 @@ -# planning-group-foundation Specification - -## Requirements - -### Requirement: Planning group definitions must identify kinematic chains -The planning configuration MUST represent kinematic chain metadata with planning-group definitions containing group-local joint names, a chain base link, and an optional pose-target tip link. - -#### Scenario: Explicit manipulator group is configured -- **WHEN** a manipulator config declares a group with local joints, `base_link`, and `tip_link` -- **THEN** the resolved robot model exposes that group without relying on a robot-scoped end-effector field - -### Requirement: Robot configuration must reject removed end-effector metadata -Robot configuration MUST NOT accept robot-scoped `end_effector_link` as the source of pose-target metadata. It MUST direct callers to planning-group `tip_link` values or SRDF discovery. - -#### Scenario: Legacy end-effector field is provided -- **WHEN** a robot config is constructed with `end_effector_link` -- **THEN** validation fails with migration guidance - -### Requirement: Group identifiers must be globally scoped -Planning-group APIs MUST use stable group identifiers scoped by robot name so multi-robot stacks can distinguish groups with the same local name. - -#### Scenario: Two robots use manipulator groups -- **WHEN** two robot configs each define a `manipulator` group -- **THEN** their resolved group IDs are distinct diff --git a/openspec/specs/planning-world-monitor-groups/spec.md b/openspec/specs/planning-world-monitor-groups/spec.md deleted file mode 100644 index 5a737525bc..0000000000 --- a/openspec/specs/planning-world-monitor-groups/spec.md +++ /dev/null @@ -1,26 +0,0 @@ -## Purpose - -Specify world and monitor support for planning-group queries. - -## Requirements - -### Requirement: World backends must answer group pose queries -World backends MUST provide end-effector pose queries for pose-targetable planning groups using the group's configured target frame. - -#### Scenario: Group pose is requested -- **WHEN** a caller requests FK for a valid pose-targetable group ID -- **THEN** the backend returns the pose of that group's `tip_link` - -### Requirement: World backends must return group-ordered Jacobians -World backends MUST return Jacobian columns ordered according to the requested group's local joint names. - -#### Scenario: Backend returns a full robot Jacobian -- **WHEN** a group contains a subset of robot joints in a custom order -- **THEN** the backend projects and reorders Jacobian columns to match the group - -### Requirement: Monitors must expose group-scoped state -World monitors MUST support planning-group state queries without requiring callers to infer robot-local joint mappings. - -#### Scenario: Group state is requested for a multi-group robot -- **WHEN** a caller requests state for a specific group ID -- **THEN** the monitor returns joint data scoped and ordered for that group From 3684ee3ee128bfac1c25f4e12526ede609aee83c Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 14 Jul 2026 18:15:46 -0700 Subject: [PATCH 21/29] feat: roboplan suport (single robot) --- .../planning/world/roboplan_world.py | 89 +++++++++++++++++-- dimos/manipulation/test_roboplan_world.py | 75 +++++++++++++++- 2 files changed, 155 insertions(+), 9 deletions(-) diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index d244bfe246..f5dd421589 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -47,6 +47,7 @@ ) from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions +from dimos.manipulation.planning.planners.selected_joint_space import normalize_selection_target from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -424,10 +425,71 @@ def plan_selected_joint_path( timeout: float = 10.0, max_iterations: int = 5000, ) -> PlanningResult: - """PlanningGroupSelection planning is not implemented by RoboPlanWorld.""" + """Plan a single planning group using RoboPlan's native RRT.""" + if world is not self: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + message="RoboPlan-native planner requires its RoboPlanWorld instance", + ) + start_time = time.time() + if not selection.groups: + return PlanningResult( + status=PlanningStatus.INVALID_GOAL, + message="No planning groups selected", + ) + if len(selection.groups) != 1: + return PlanningResult( + status=PlanningStatus.UNSUPPORTED, + message="RoboPlan-native planning supports exactly one selected planning group", + ) + + group = selection.groups[0] + try: + normalized_start = normalize_selection_target(selection, start, "start") + except ValueError as exc: + return PlanningResult(status=PlanningStatus.INVALID_START, message=str(exc)) + try: + normalized_goal = normalize_selection_target(selection, goal, "goal") + except ValueError as exc: + return PlanningResult(status=PlanningStatus.INVALID_GOAL, message=str(exc)) + + robot_id = self._robot_id_for_group(group.id) + q_start = np.asarray(normalized_start.position, dtype=np.float64) + q_goal = np.asarray(normalized_goal.position, dtype=np.float64) + try: + path_arrays = self._run_native_rrt( + robot_id, + q_start, + q_goal, + timeout, + group_name=group.group_name, + joint_names=list(group.local_joint_names), + ) + except ValueError as exc: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + planning_time=time.time() - start_time, + message=f"RoboPlan-native planning failed: {exc}", + ) + if not path_arrays: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + planning_time=time.time() - start_time, + message="RoboPlan-native planning failed: returned an empty path", + ) + path = [ + JointState( + name=list(selection.joint_names), + position=np.asarray(q).astype(float).tolist(), + ) + for q in path_arrays + ] return PlanningResult( - status=PlanningStatus.UNSUPPORTED, - message="RoboPlanWorld does not support selected planning-group joint paths", + status=PlanningStatus.SUCCESS, + path=path, + planning_time=time.time() - start_time, + path_length=compute_path_length(path), + message="RoboPlan path found", ) def get_name(self) -> str: @@ -475,6 +537,11 @@ def _generate_srdf(self, config: RobotModelConfig, urdf_path: Path) -> str: for joint_name in config.joint_names: lines.append(f' ') lines.append(" ") + for group in config.planning_groups: + lines.append(f' ') + for joint_name in group.joint_names: + lines.append(f' ') + lines.append(" ") for link1, link2 in self._collision_exclusion_pairs(config, urdf_path): lines.append( f' list[NDArray[np.float64]]: scene = self._require_scene() robot = self._get_robot(robot_id) options = roboplan_rrt.RRTOptions() - options.group_name = robot.config.name + native_group_name = robot.config.name if group_name is None else group_name + native_joint_names = robot.config.joint_names if joint_names is None else joint_names + options.group_name = native_group_name options.max_planning_time = timeout options.collision_check_use_bisection = False planner = roboplan_rrt.RRT(scene, options) - start_config = self._to_native_joint_configuration(robot_id, q_start) - goal_config = self._to_native_joint_configuration(robot_id, q_goal) + start_config = self._to_native_joint_configuration(robot_id, q_start, native_joint_names) + goal_config = self._to_native_joint_configuration(robot_id, q_goal, native_joint_names) result = planner.plan(start_config, goal_config) if result is None: raise ValueError("RoboPlan RRT returned no path") return self._extract_native_path(result) def _to_native_joint_configuration( - self, robot_id: WorldRobotID, q: NDArray[np.float64] + self, robot_id: WorldRobotID, q: NDArray[np.float64], joint_names: list[str] | None = None ) -> roboplan_core.JointConfiguration: robot = self._get_robot(robot_id) return roboplan_core.JointConfiguration( - robot.config.joint_names, np.asarray(q, dtype=np.float64) + robot.config.joint_names if joint_names is None else joint_names, + np.asarray(q, dtype=np.float64), ) def _extract_native_path(self, result: roboplan_core.JointPath) -> list[NDArray[np.float64]]: diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 32eaa97bfa..5f7612a693 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -25,7 +25,10 @@ import numpy as np import pytest -from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.models import ( + PlanningGroupDefinition, + PlanningGroupSelection, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import Obstacle @@ -732,6 +735,76 @@ 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_selected_planner_returns_global_selected_joint_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + group = world._planning_group_from_id("arm/manipulator") + selection = PlanningGroupSelection.from_groups((group,)) + + result = world.plan_selected_joint_path( + world, + selection, + JointState(name=["arm/joint1", "arm/joint2"], position=[0.0, 0.0]), + JointState(name=["arm/joint1", "arm/joint2"], position=[0.4, 0.2]), + timeout=1.0, + ) + + assert result.status == PlanningStatus.SUCCESS + assert [state.name for state in result.path] == [["arm/joint1", "arm/joint2"]] * 3 + assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + + +def test_native_selected_planner_accepts_local_joint_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + selection = PlanningGroupSelection.from_groups( + (world._planning_group_from_id("arm/manipulator"),) + ) + + result = world.plan_selected_joint_path( + world, + selection, + JointState(name=["joint2", "joint1"], position=[0.2, 0.0]), + JointState(name=["joint2", "joint1"], position=[0.4, 0.2]), + ) + + assert result.status == PlanningStatus.SUCCESS + assert result.path[0].name == ["arm/joint1", "arm/joint2"] + assert result.path[0].position == [0.0, 0.2] + + +def test_native_selected_planner_rejects_multi_group_selection( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition("left", ("joint1",), "base", "left_tip"), + PlanningGroupDefinition("right", ("joint2",), "base", "right_tip"), + ] + } + ) + world, _ = _make_world(fake_roboplan, config) + world.finalize() + selection = PlanningGroupSelection.from_groups( + (world._planning_group_from_id("arm/left"), world._planning_group_from_id("arm/right")) + ) + + 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.1, 0.1]), + ) + + assert result.status == PlanningStatus.UNSUPPORTED + assert "exactly one" in result.message + + def test_native_planner_rejects_empty_path( fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch ) -> None: From 9859e03cf79a46827055da5b0de370879ffae2fe Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 14 Jul 2026 18:50:35 -0700 Subject: [PATCH 22/29] doc: alignment --- .../manipulation/planning_groups.md | 39 +++++++++++-------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/docs/capabilities/manipulation/planning_groups.md b/docs/capabilities/manipulation/planning_groups.md index f2dac09afb..7ce9edb122 100644 --- a/docs/capabilities/manipulation/planning_groups.md +++ b/docs/capabilities/manipulation/planning_groups.md @@ -23,17 +23,19 @@ use global joint names so multiple robots can safely share local names such as ## Discovering planning groups -Planning-group config helpers can discover groups in this order: +Robot configs can provide planning groups explicitly with +`RobotModelConfig.planning_groups`. Direct `RobotModelConfig(...)` construction +does not run discovery or synthesize groups in `model_post_init`; callers must +pass explicit `planning_groups` there. -1. Explicit `planning_groups` on the robot model config. -2. Explicit `srdf_path` on the robot model config. -3. Conservative SRDF auto-discovery near the model path, with a warning. -4. Fallback generation of one `{robot_name}/manipulator` group when the - configured controllable joints form exactly one unambiguous serial chain. -5. Error if no SRDF or fallback chain can provide a single valid group. +When code uses the discovery helper instead of explicit config, DimOS discovers +groups in this order: -Direct `RobotModelConfig(...)` construction does not run discovery or synthesize -groups in `model_post_init`; callers must pass explicit `planning_groups` there. +1. Explicit `srdf_path` provided to the helper. +2. Conservative SRDF auto-discovery near the model path, with a warning. +3. Fallback generation of one `{robot_name}/manipulator` group when the + configured controllable joints form exactly one unambiguous serial chain. +4. Error if no SRDF or fallback chain can provide a single valid group. Supported SRDF group forms: @@ -59,8 +61,8 @@ unique serial target frame. ## Fallback behavior -When no SRDF or explicit group config is available, -fallback uses `RobotModelConfig.joint_names` as the candidate controllable set. +When discovery runs without an SRDF, fallback uses +`RobotModelConfig.joint_names` as the candidate controllable set. This field is the robot's ordered local model joint set, not an implicit planning group. @@ -80,8 +82,10 @@ pose_groups = [group for group in groups if group.has_pose_target] group_id = pose_groups[0].id ``` -Joint-space planning targets group IDs and uses local joint names inside each -target `JointState`: +Joint-space planning targets group IDs. Each target `JointState` may be +unnamed in the group's joint order, named with all local model joint names, or +named with all global joint names. Do not mix local and global names in one +target. ```python skip ok = manip.plan_to_joint_targets( @@ -95,7 +99,8 @@ ok = manip.plan_to_joint_targets( ``` Pose planning targets pose-capable group IDs. Add auxiliary groups when another -chain should participate as free DOFs but does not have its own pose target: +chain should participate as free DOFs but does not have its own pose target. +Pose targets are `Pose` values keyed by planning group ID: ```python skip ok = manip.plan_to_pose_targets( @@ -120,9 +125,9 @@ manip.execute_plan(plan) ``` For robot-scoped compatibility APIs, unnamed joint vectors are interpreted in -the robot model's configured joint order. If names are provided, they must be -local model joint names: no global names, missing joints, extra joints, or -partial joint sets. +the selected default planning group's joint order. If names are provided, they +may be all local model joint names or all global joint names. Missing joints, +extra joints, partial joint sets, and mixed local/global namespaces are rejected. ## Generated plans and execution From 35397cc9db2b2615aa0a19d6605f12a6ccb0a2a8 Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 14 Jul 2026 23:10:39 -0700 Subject: [PATCH 23/29] fix: underscore ass --- dimos/manipulation/manipulation_module.py | 17 ++++++++--------- .../planning/groups/test_planning_groups.py | 4 ++-- .../planning/monitor/test_world_monitor.py | 1 - .../manipulation/planning/world/drake_world.py | 2 -- dimos/manipulation/test_roboplan_world.py | 3 --- dimos/manipulation/visualization/factory.py | 1 - .../manipulation/visualization/test_operator.py | 5 ----- dimos/manipulation/visualization/viser/gui.py | 1 - .../visualization/viser/test_gui.py | 2 -- .../viser/test_viser_visualization.py | 2 -- 10 files changed, 10 insertions(+), 28 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 4db48588d4..f1db951cb6 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -705,14 +705,14 @@ def _plan_selected_path( timeout=self.config.planning_timeout, ) if not result.is_success(): - _ = self._fail_planning_epoch(planning_epoch, f"Planning failed: {result.status.name}") + self._fail_planning_epoch(planning_epoch, f"Planning failed: {result.status.name}") return None logger.info("Path: %d waypoints, groups=%s", len(result.path), group_ids) try: path, trajectory = self._materialize_generated_plan(group_ids, result.path) except Exception as exc: - _ = self._fail_planning_epoch(planning_epoch, f"Failed to materialize plan: {exc}") + self._fail_planning_epoch(planning_epoch, f"Failed to materialize plan: {exc}") return None plan = GeneratedPlan( group_ids=group_ids, @@ -967,7 +967,7 @@ def generate_plan_to_joint_targets( if self._world_monitor is None or self._planner is None: return None if not joint_targets: - _ = self._fail("At least one joint target is required") + self._fail("At least one joint target is required") return None group_ids = tuple( @@ -982,7 +982,7 @@ def generate_plan_to_joint_targets( current = self._world_monitor.current_global_joint_state() start = filter_joint_state_to_selected_joints(current, selection.joint_names) except Exception as exc: - _ = self._fail_planning_epoch( + self._fail_planning_epoch( planning_epoch, f"Failed to resolve planning groups: {exc}" ) return None @@ -996,7 +996,7 @@ def generate_plan_to_joint_targets( target_global = joint_target_to_global_names(target_group, target) except (KeyError, ValueError) as exc: logger.error(str(exc)) - _ = self._fail_planning_epoch( + self._fail_planning_epoch( planning_epoch, f"Invalid joint target for '{group_id}'" ) return None @@ -1015,7 +1015,7 @@ def generate_plan_to_pose_targets( if self._world_monitor is None or self._kinematics is None: return None if not pose_targets: - _ = self._fail("At least one pose target is required") + self._fail("At least one pose target is required") return None stamped_targets = { planning_group_id_from_selector(group): PoseStamped( @@ -1035,7 +1035,7 @@ def generate_plan_to_pose_targets( current = self._world_monitor.current_global_joint_state() start = filter_joint_state_to_selected_joints(current, selection.joint_names) except Exception as exc: - _ = self._fail_planning_epoch( + self._fail_planning_epoch( planning_epoch, f"Failed to resolve planning groups: {exc}" ) return None @@ -1045,7 +1045,7 @@ def generate_plan_to_pose_targets( seed=start, ) if not ik.is_success() or ik.joint_state is None: - _ = self._fail_planning_epoch(planning_epoch, f"IK failed: {ik.status.name}") + self._fail_planning_epoch(planning_epoch, f"IK failed: {ik.status.name}") return None logger.info(f"IK solved, error: {ik.position_error:.4f}m") return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) @@ -1091,7 +1091,6 @@ def preview_plan( return False if self._world_monitor is None: return False - _ = target_fps self._world_monitor.animate_trajectory(plan.trajectory, duration) return True diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 5c18011a62..bcc5d0b17a 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -627,7 +627,7 @@ def test_robot_model_config_end_effector_link_requires_pose_group() -> None: ) with pytest.raises(ValueError, match="no pose-target planning group"): - _ = config.end_effector_link + getattr(config, "end_effector_link") def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> None: @@ -652,7 +652,7 @@ def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> ) with pytest.raises(ValueError, match="multiple pose-target planning groups"): - _ = config.end_effector_link + getattr(config, "end_effector_link") def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index a9781a6c8a..aa868e2892 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -54,7 +54,6 @@ def get_current_velocities(self) -> None: return None def is_state_stale(self, max_age: float) -> bool: - _ = max_age return self._stale diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 09bb30dc80..3474990e35 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -1157,7 +1157,6 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray def initialize(self, session: VisualizationSession) -> None: """Embedded Meshcat observes the Drake world directly; no extra sync needed.""" - _ = session return None def get_visualization_url(self) -> str | None: @@ -1178,7 +1177,6 @@ def _publish_visualization(self, ctx: Context | None = None) -> None: def update_state(self, frame: VisualizationStateFrame) -> None: """Receive pushed state frame; embedded Meshcat uses Drake live context.""" - _ = frame self._publish_visualization() def _set_preview_positions( diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 186d75638f..f24d010ecc 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -521,7 +521,6 @@ def test_group_fk_and_jacobian_use_group_tip_and_local_joint_order( def fake_fk( self: FakeScene, q: np.ndarray, frame_name: str, base_frame: str = "" ) -> np.ndarray: - _ = (self, base_frame) fk_frames.append(frame_name) mat = np.eye(4) mat[0, 3] = float(np.sum(q)) @@ -530,7 +529,6 @@ def fake_fk( def fake_jacobian( self: FakeScene, q: np.ndarray, frame_name: str, local: bool = True ) -> np.ndarray: - _ = (self, q) assert frame_name == "wrist_tip" assert local is True return np.arange(18, dtype=np.float64).reshape(6, 3) @@ -635,7 +633,6 @@ def test_group_jacobian_falls_back_to_configured_joint_order_when_scene_order_is world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) def missing_group_info(self: FakeScene, name: str) -> FakeJointGroupInfo: - _ = (self, name) raise AttributeError("no joint group info") monkeypatch.setattr(FakeScene, "getJointGroupInfo", missing_group_info) diff --git a/dimos/manipulation/visualization/factory.py b/dimos/manipulation/visualization/factory.py index cbb2fb397d..3d9d9a230d 100644 --- a/dimos/manipulation/visualization/factory.py +++ b/dimos/manipulation/visualization/factory.py @@ -53,7 +53,6 @@ def create_manipulation_visualization( ViserManipulationVisualizer, ) - _ = world_monitor, manipulation_module return ViserManipulationVisualizer(config=config) raise AssertionError(f"Unhandled manipulation visualization config: {config!r}") diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py index 2c15ef9326..701e6625da 100644 --- a/dimos/manipulation/visualization/test_operator.py +++ b/dimos/manipulation/visualization/test_operator.py @@ -172,12 +172,9 @@ def generate_plan_to_pose_targets( def preview_plan( self, plan: GeneratedPlan | None = None, duration: float | None = None ) -> bool: - _ = plan - _ = duration return self.preview_success def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: - _ = plan return self.execute_success def cancel(self) -> bool: @@ -207,13 +204,11 @@ def get_current_joint_state(self, robot_id: str) -> JointState | None: return self.current_states.get(robot_id) def is_state_valid(self, robot_id: str, joint_state: JointState) -> bool: - _ = robot_id, joint_state return self.valid def get_group_ee_pose( self, group_id: PlanningGroupID, joint_state: JointState | None = None ) -> PoseStamped: - _ = group_id, joint_state return PoseStamped( frame_id="world", position=Vector3(1.0, 2.0, 3.0), orientation=Quaternion() ) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index 56eb548ee9..e959f82f76 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -226,7 +226,6 @@ def _current_target_for_group(self, group: PlanningGroup) -> dict[PlanningGroupI } def is_state_stale(self, robot_name: RobotName, max_age: float = 1.0) -> bool: - _ = max_age return self.get_current_joint_state(robot_name) is None def get_module_state(self) -> str: diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index e4e83f6448..9d8be1e168 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -61,7 +61,6 @@ def status(self) -> OperatorStatus: return OperatorStatus(state="IDLE", error="", has_plan=False) def get_init_joints(self, robot_name: str) -> None: - _ = robot_name return None def cancel(self) -> bool: @@ -135,7 +134,6 @@ def submit( timeout_seconds: float | None = None, on_error: Callable[[str], None] | None = None, ) -> None: - _ = operation, timeout_seconds if on_error is not None: self.errors.append(on_error) diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 6a1d85e432..c4c638374f 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -369,11 +369,9 @@ def plan_to_pose(self, request: object) -> GeneratedPlan: return self.module.make_plan(tuple(request.pose_targets)) # type: ignore[attr-defined] def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: - _ = plan, duration return self.module.preview_plan() def execute(self, plan: GeneratedPlan) -> bool: - _ = plan return self.module.execute() def cancel(self) -> bool: From 72e71b38b530be19a957025238061b6c950809b3 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 06:11:11 +0000 Subject: [PATCH 24/29] [autofix.ci] apply automated fixes --- dimos/manipulation/manipulation_module.py | 12 +++--------- .../planning/groups/test_planning_groups.py | 4 ++-- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index f1db951cb6..1efdf35c6c 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -982,9 +982,7 @@ def generate_plan_to_joint_targets( current = self._world_monitor.current_global_joint_state() start = filter_joint_state_to_selected_joints(current, selection.joint_names) except Exception as exc: - self._fail_planning_epoch( - planning_epoch, f"Failed to resolve planning groups: {exc}" - ) + self._fail_planning_epoch(planning_epoch, f"Failed to resolve planning groups: {exc}") return None goal_names: list[str] = [] @@ -996,9 +994,7 @@ def generate_plan_to_joint_targets( target_global = joint_target_to_global_names(target_group, target) except (KeyError, ValueError) as exc: logger.error(str(exc)) - self._fail_planning_epoch( - planning_epoch, f"Invalid joint target for '{group_id}'" - ) + self._fail_planning_epoch(planning_epoch, f"Invalid joint target for '{group_id}'") return None goal_names.extend(target_global.name) goal_positions.extend(target_global.position) @@ -1035,9 +1031,7 @@ def generate_plan_to_pose_targets( current = self._world_monitor.current_global_joint_state() start = filter_joint_state_to_selected_joints(current, selection.joint_names) except Exception as exc: - self._fail_planning_epoch( - planning_epoch, f"Failed to resolve planning groups: {exc}" - ) + self._fail_planning_epoch(planning_epoch, f"Failed to resolve planning groups: {exc}") return None ik = self.inverse_kinematics( pose_targets=stamped_targets, diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index bcc5d0b17a..44ac55b9b9 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -627,7 +627,7 @@ def test_robot_model_config_end_effector_link_requires_pose_group() -> None: ) with pytest.raises(ValueError, match="no pose-target planning group"): - getattr(config, "end_effector_link") + config.end_effector_link def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> None: @@ -652,7 +652,7 @@ def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> ) with pytest.raises(ValueError, match="multiple pose-target planning groups"): - getattr(config, "end_effector_link") + config.end_effector_link def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: From f708240d20ae6e969aa6618863ea5c834be73361 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 15 Jul 2026 20:37:05 -0700 Subject: [PATCH 25/29] test: satisfy ruff in planning group tests --- dimos/manipulation/planning/groups/test_planning_groups.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 44ac55b9b9..88757ec254 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -627,7 +627,7 @@ def test_robot_model_config_end_effector_link_requires_pose_group() -> None: ) with pytest.raises(ValueError, match="no pose-target planning group"): - config.end_effector_link + assert config.end_effector_link def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> None: @@ -652,7 +652,7 @@ def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> ) with pytest.raises(ValueError, match="multiple pose-target planning groups"): - config.end_effector_link + assert config.end_effector_link def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: From 50806573c687011efafea9de8023ccc5de6adf37 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 15 Jul 2026 21:21:25 -0700 Subject: [PATCH 26/29] fix: make manipulation execution cancellation safe --- dimos/manipulation/manipulation_module.py | 193 ++++++++++++++---- dimos/manipulation/test_manipulation_unit.py | 14 ++ .../test_plan_execution_reservation.py | 45 +++- 3 files changed, 211 insertions(+), 41 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 1efdf35c6c..08d546e19a 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -155,6 +155,13 @@ def __init__(self, **kwargs: Any) -> None: # State machine self._state = ManipulationState.IDLE self._lock = threading.Lock() + # State transitions are local, but coordinator calls must be serialized + # with cancellation. In particular, do not let cancel race between + # two arms of a multi-arm dispatch. + self._execution_dispatch_lock = threading.Lock() + self._execution_generation = 0 + self._possibly_active_tasks: set[str] = set() + self._cancellation_in_progress = False self._error_message = "" self._planning_epoch = 0 @@ -182,6 +189,17 @@ def __init__(self, **kwargs: Any) -> None: logger.info("ManipulationModule initialized") + def _ensure_execution_transaction_state(self) -> None: + """Initialize transaction fields for lightweight test harness subclasses.""" + if not hasattr(self, "_execution_dispatch_lock"): + self._execution_dispatch_lock = threading.Lock() + if not hasattr(self, "_execution_generation"): + self._execution_generation = 0 + if not hasattr(self, "_possibly_active_tasks"): + self._possibly_active_tasks = set() + if not hasattr(self, "_cancellation_in_progress"): + self._cancellation_in_progress = False + @rpc def start(self) -> None: """Start the manipulation module.""" @@ -407,15 +425,43 @@ def get_error(self) -> str: @rpc def cancel(self) -> bool: """Cancel current motion or invalidate an in-progress plan.""" + self._ensure_execution_transaction_state() + # Publish the cancellation gate before waiting for dispatch. An + # execution that arrives while cancellation is in flight must not + # become a new transaction behind the cancellation barrier. with self._lock: - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - elif self._state == ManipulationState.EXECUTING: - self._state = ManipulationState.IDLE - else: - return False - plan = self._last_plan + self._cancellation_in_progress = True + # Holding this lock waits for an in-flight execute RPC to resolve and + # prevents another arm from being dispatched after cancellation. + try: + with self._execution_dispatch_lock: + with self._lock: + task_names = set(self._possibly_active_tasks) + is_planning = self._state == ManipulationState.PLANNING + is_executing = self._state == ManipulationState.EXECUTING + if not (is_planning or is_executing or task_names): + return False + self._execution_generation += 1 + if is_planning: + self._planning_epoch += 1 + plan = self._last_plan + + unresolved = self._cancel_tasks(task_names) + with self._lock: + if unresolved: + self._possibly_active_tasks = unresolved + self._state = ManipulationState.FAULT + self._error_message = ( + "Failed to cancel coordinator tasks: " + + ", ".join(sorted(unresolved)) + ) + logger.error(self._error_message) + return False + self._possibly_active_tasks.clear() + self._state = ManipulationState.IDLE + finally: + with self._lock: + self._cancellation_in_progress = False if plan is not None: self._dismiss_preview(plan.group_ids) logger.info("Motion cancelled") @@ -1429,26 +1475,30 @@ def execute_plan( execution = self._begin_execution(plan) if execution is None: return False - previous_state, target_plan = execution + previous_state, target_plan, token = execution try: freshness_error = self._stored_plan_freshness_error(target_plan) if freshness_error is not None: - self._restore_execution_gate(previous_state, freshness_error) + self._restore_execution_gate(token, previous_state, freshness_error) return False prepared = self._prepare_execution(target_plan, robot_name) if isinstance(prepared, str): - self._restore_execution_gate(previous_state, prepared) + self._restore_execution_gate(token, previous_state, prepared) return False except Exception as exc: - self._restore_execution_gate(previous_state, f"Failed to prepare execution: {exc}") + self._restore_execution_gate(token, previous_state, f"Failed to prepare execution: {exc}") return False - return self._dispatch_prepared_plan(prepared) + return self._dispatch_prepared_plan(prepared, token) def _begin_execution( self, plan: GeneratedPlan | None - ) -> tuple[ManipulationState, GeneratedPlan] | None: + ) -> tuple[ManipulationState, GeneratedPlan, int] | None: """Enter EXECUTING if the module is idle enough to dispatch.""" + self._ensure_execution_transaction_state() with self._lock: + if self._cancellation_in_progress: + logger.warning("Cancellation is in progress; rejecting execution") + return None target_plan = plan or self._last_plan if target_plan is None or not target_plan.path or not target_plan.trajectory.points: logger.warning("Stored plan is invalid or not executable") @@ -1458,14 +1508,19 @@ def _begin_execution( logger.warning("Manipulation state is not executable") return None previous_state = self._state + self._execution_generation += 1 + token = self._execution_generation self._state = ManipulationState.EXECUTING - return previous_state, target_plan + return previous_state, target_plan, token - def _restore_execution_gate(self, previous_state: ManipulationState, msg: str) -> None: + def _restore_execution_gate( + self, token: int, previous_state: ManipulationState, msg: str + ) -> None: logger.warning(msg) with self._lock: - self._state = previous_state - self._error_message = msg + if self._execution_generation == token and self._state == ManipulationState.EXECUTING: + self._state = previous_state + self._error_message = msg def _stored_plan_freshness_error( self, plan: GeneratedPlan, tolerance: float = 1e-6 @@ -1552,34 +1607,92 @@ def _prepare_execution( return tuple(prepared) def _dispatch_prepared_plan( - self, prepared: Sequence[tuple[RobotName, RobotModelConfig, JointTrajectory]] + self, + prepared: Sequence[tuple[RobotName, RobotModelConfig, JointTrajectory]], + token: int, ) -> bool: """Dispatch already-validated robot trajectories; failures here fault.""" - if (client := self._get_coordinator_client()) is None: - logger.error("No coordinator client") - return False + self._ensure_execution_transaction_state() + with self._execution_dispatch_lock: + if (client := self._get_coordinator_client()) is None: + return self._fault_execution(token, "No coordinator client") + + for _robot_name, config, traj in prepared: + task_name = config.coordinator_task_name + if task_name is None: + return self._fault_execution(token, "No coordinator_task_name") + with self._lock: + if ( + self._execution_generation != token + or self._state != ManipulationState.EXECUTING + ): + return False + # Add before the RPC: a rejected or raised execute may + # have been accepted remotely. + self._possibly_active_tasks.add(task_name) + try: + translated = self._translate_trajectory_to_coordinator(traj, config) + logger.info( + "Executing: task='%s', %d pts, %.2fs", + task_name, + len(translated.points), + translated.duration, + ) + result = client.task_invoke(task_name, "execute", {"trajectory": translated}) + except Exception as exc: + return self._fault_execution(token, f"Failed to dispatch trajectory: {exc}") + if not result: + return self._fault_execution(token, "Coordinator rejected trajectory") - for _robot_name, config, traj in prepared: + logger.info("Trajectory accepted") + with self._lock: + if ( + self._execution_generation != token + or self._state != ManipulationState.EXECUTING + ): + return False + self._state = ManipulationState.COMPLETED + return True + + def _cancel_tasks(self, task_names: set[str]) -> set[str]: + """Cancel possibly active tasks, returning names whose safety is unknown.""" + if not task_names: + return set() + client = self._get_coordinator_client() + if client is None: + return set(task_names) + unresolved: set[str] = set() + for task_name in task_names: try: - translated = self._translate_trajectory_to_coordinator(traj, config) - logger.info( - "Executing: task='%s', %d pts, %.2fs", - config.coordinator_task_name, - len(translated.points), - translated.duration, - ) - result = client.task_invoke( - config.coordinator_task_name, "execute", {"trajectory": translated} - ) - except Exception as exc: - return self._fail(f"Failed to dispatch trajectory: {exc}") - if not result: - return self._fail("Coordinator rejected trajectory") + result = client.task_invoke(task_name, "cancel", {}) + # Coordinator cancel returns False when the task is already + # non-executing; that is a confirmed-safe result. + if result is None: + unresolved.add(task_name) + except Exception: + unresolved.add(task_name) + return unresolved - logger.info("Trajectory accepted") + def _fault_execution(self, token: int, message: str) -> bool: + """Rollback attempted tasks and fault, without overwriting cancellation.""" + self._ensure_execution_transaction_state() with self._lock: - self._state = ManipulationState.COMPLETED - return True + if self._execution_generation != token: + return False + task_names = set(self._possibly_active_tasks) + unresolved = self._cancel_tasks(task_names) + with self._lock: + if self._execution_generation != token: + return False + if unresolved: + self._possibly_active_tasks = unresolved + message += "; unresolved coordinator tasks: " + ", ".join(sorted(unresolved)) + else: + self._possibly_active_tasks.clear() + self._state = ManipulationState.FAULT + self._error_message = message + logger.error(message) + return False @rpc def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 4f37c28fa0..4ad998517e 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -228,6 +228,20 @@ def test_cancel_hides_active_plan_preview(self): module._world_monitor.cancel_preview_animation.assert_called_once_with() + def test_cancel_completed_execution_cancels_coordinator_task(self): + module = _make_module() + module._state = ManipulationState.COMPLETED + module._possibly_active_tasks = {"traj_arm"} + module._robots = { + "arm": ("arm_id", MagicMock(coordinator_task_name="traj_arm"), MagicMock()) + } + module._coordinator_client = MagicMock() + module._coordinator_client.task_invoke.return_value = False + + assert module.cancel() is True + module._coordinator_client.task_invoke.assert_called_once_with("traj_arm", "cancel", {}) + assert module._state == ManipulationState.IDLE + def test_reset_not_during_execution(self): """Reset works in any state except EXECUTING.""" module = _make_module() diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py index 74eee2cd1a..1977c6a94d 100644 --- a/dimos/manipulation/test_plan_execution_reservation.py +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -206,6 +206,48 @@ def call_execute(execute: Callable[[], bool]) -> None: assert module._coordinator_client.task_invoke.call_count == 1 +def test_execute_started_during_cancel_cannot_dispatch_after_cancel(): + module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) + module._last_plan = _plan() + module._state = ManipulationState.COMPLETED + module._possibly_active_tasks = {"traj_arm"} + cancel_started = threading.Event() + execute_finished = threading.Event() + release_cancel = threading.Event() + cancel_result: list[bool] = [] + execute_result: list[bool] = [] + + def invoke(task_name: str, method: str, _payload: dict[str, object]) -> bool: + assert task_name == "traj_arm" + if method == "cancel": + cancel_started.set() + assert release_cancel.wait(timeout=1.0) + return False + assert method == "execute" + raise AssertionError("execution dispatched during cancellation") + + module._coordinator_client.task_invoke.side_effect = invoke + cancelling = threading.Thread(target=lambda: cancel_result.append(module.cancel())) + cancelling.start() + assert cancel_started.wait(timeout=1.0) + + def run_execute() -> None: + execute_result.append(module.execute_plan()) + execute_finished.set() + + executing = threading.Thread(target=run_execute) + executing.start() + assert execute_finished.wait(timeout=1.0) + assert execute_result == [False] + assert not any(call.args[1] == "execute" for call in module._coordinator_client.task_invoke.call_args_list) + release_cancel.set() + cancelling.join(timeout=1.0) + executing.join(timeout=1.0) + + assert cancel_result == [True] + module._coordinator_client.task_invoke.assert_called_once_with("traj_arm", "cancel", {}) + + def test_execute_plan_accepts_a_direct_plan_without_reserving_stored_plan(): module = make_module() direct_plan = _plan() @@ -268,4 +310,5 @@ def test_execute_plan_dispatch_exception_faults_without_sticking_executing(): assert module.execute_plan() is False assert module._state == ManipulationState.FAULT - assert module._error_message == "Failed to dispatch trajectory: rpc exploded" + assert module._error_message.startswith("Failed to dispatch trajectory: rpc exploded") + assert "traj_arm" in module._error_message From c64946c0fd76fd443b176477ac4b5f4e412a4eb3 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 04:21:51 +0000 Subject: [PATCH 27/29] [autofix.ci] apply automated fixes --- dimos/manipulation/manipulation_module.py | 9 +++++---- dimos/manipulation/test_plan_execution_reservation.py | 4 +++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 08d546e19a..1bb55dbc89 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -451,9 +451,8 @@ def cancel(self) -> bool: if unresolved: self._possibly_active_tasks = unresolved self._state = ManipulationState.FAULT - self._error_message = ( - "Failed to cancel coordinator tasks: " - + ", ".join(sorted(unresolved)) + self._error_message = "Failed to cancel coordinator tasks: " + ", ".join( + sorted(unresolved) ) logger.error(self._error_message) return False @@ -1486,7 +1485,9 @@ def execute_plan( self._restore_execution_gate(token, previous_state, prepared) return False except Exception as exc: - self._restore_execution_gate(token, previous_state, f"Failed to prepare execution: {exc}") + self._restore_execution_gate( + token, previous_state, f"Failed to prepare execution: {exc}" + ) return False return self._dispatch_prepared_plan(prepared, token) diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py index 1977c6a94d..935804834d 100644 --- a/dimos/manipulation/test_plan_execution_reservation.py +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -239,7 +239,9 @@ def run_execute() -> None: executing.start() assert execute_finished.wait(timeout=1.0) assert execute_result == [False] - assert not any(call.args[1] == "execute" for call in module._coordinator_client.task_invoke.call_args_list) + assert not any( + call.args[1] == "execute" for call in module._coordinator_client.task_invoke.call_args_list + ) release_cancel.set() cancelling.join(timeout=1.0) executing.join(timeout=1.0) From ea383bff9b0be7df2f76c354f37fe8b270dea25d Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 21 Jul 2026 00:42:36 -0700 Subject: [PATCH 28/29] refactor(manipulation): extract execution runtime --- .../_test_manipulation_helpers.py | 86 +- .../control/coordinator_client.py | 16 +- dimos/manipulation/execution_runtime.py | 2084 +++++++++++++++++ dimos/manipulation/manipulation_module.py | 1204 +++++----- .../planning/examples/manipulation_client.py | 8 +- dimos/manipulation/skill_errors.py | 6 + dimos/manipulation/test_execution_runtime.py | 1811 ++++++++++++++ .../test_generated_plan_materialization.py | 59 +- .../manipulation/test_manipulation_module.py | 103 +- .../test_manipulation_monitor_preview.py | 243 +- dimos/manipulation/test_manipulation_unit.py | 734 +----- .../test_plan_execution_reservation.py | 414 ++-- dimos/manipulation/visualization/operator.py | 29 +- .../visualization/test_operator.py | 34 +- dimos/manipulation/visualization/viser/gui.py | 90 +- .../manipulation/visualization/viser/state.py | 25 +- .../visualization/viser/test_gui.py | 62 +- .../viser/test_viser_visualization.py | 17 +- .../.openspec.yaml | 2 + .../design.md | 115 + .../proposal.md | 29 + .../coordinator-trajectory-execution/spec.md | 85 + .../manipulation-lifecycle-runtime/spec.md | 96 + .../tasks.md | 71 + 24 files changed, 5447 insertions(+), 1976 deletions(-) create mode 100644 dimos/manipulation/execution_runtime.py create mode 100644 dimos/manipulation/test_execution_runtime.py create mode 100644 openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml create mode 100644 openspec/changes/extract-manipulation-execution-runtime/design.md create mode 100644 openspec/changes/extract-manipulation-execution-runtime/proposal.md create mode 100644 openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md create mode 100644 openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md create mode 100644 openspec/changes/extract-manipulation-execution-runtime/tasks.md diff --git a/dimos/manipulation/_test_manipulation_helpers.py b/dimos/manipulation/_test_manipulation_helpers.py index 6df020d7bc..35d089cb58 100644 --- a/dimos/manipulation/_test_manipulation_helpers.py +++ b/dimos/manipulation/_test_manipulation_helpers.py @@ -14,32 +14,96 @@ """Shared lightweight test harnesses for manipulation module tests.""" -import threading from unittest.mock import MagicMock -from dimos.manipulation.manipulation_module import ( - ManipulationModule, - ManipulationState, -) +from dimos.manipulation.execution_runtime import ExecutionRuntime, Outcome +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.spec.config import RobotModelConfig + +_TEST_RUNTIMES: list[ExecutionRuntime] = [] + + +class FakeCoordinatorGateway: + """Deterministic coordinator gateway used by module-side tests.""" + + def __init__(self, execute_outcome: Outcome = Outcome.ACCEPTED) -> None: + self.execute_outcome = execute_outcome + self.execute_calls: list[tuple[str, object]] = [] + self.cancel_calls: list[str] = [] + self.stopped = False + + def execute(self, task_name: str, request: object) -> Outcome: + self.execute_calls.append((task_name, request)) + return self.execute_outcome + + def cancel(self, task_name: str) -> Outcome: + self.cancel_calls.append(task_name) + return Outcome.CANCELLED + + def status(self, task_name: str) -> Outcome: + del task_name + return Outcome.INACTIVE + + def reset(self, task_name: str) -> Outcome: + del task_name + return Outcome.INACTIVE + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + del hardware_id, position + return Outcome.ACCEPTED + + def get_gripper_position(self, hardware_id: str) -> float | None: + del hardware_id + return 0.0 + + def stop(self) -> None: + self.stopped = True class ManipulationModuleHarness(ManipulationModule): """Manipulation module initialized only with state needed by unit tests.""" def __init__(self) -> None: - self._state = ManipulationState.IDLE - self._lock = threading.Lock() - self._error_message = "" - self._planning_epoch = 0 self._robots = {} - self._last_plan = None self._world_monitor = None self._planner = None self._kinematics = None - self._coordinator_client = None + self._execution_runtime = None + self._execution_topology = None + self._init_joints = {} self.config = MagicMock(planning_timeout=10.0) def make_module() -> ManipulationModule: """Create a lightweight ManipulationModule harness for behavior tests.""" return ManipulationModuleHarness() + + +def install_runtime( + module: ManipulationModule, + robots: list[RobotModelConfig], + gateway: FakeCoordinatorGateway | None = None, +) -> FakeCoordinatorGateway: + """Install the production runtime with a test-owned fake gateway.""" + gateway = gateway or FakeCoordinatorGateway() + from dimos.manipulation.execution_runtime import ExecutionTopology + + module._execution_topology = ExecutionTopology.from_robot_configs(robots) + module._execution_runtime = ExecutionRuntime( + lambda: gateway, + topology=module._execution_topology, + action_timeout=module.config.planning_timeout, + poll_interval=0.01, + ) + _TEST_RUNTIMES.append(module._execution_runtime) + return gateway + + +def close_test_runtimes() -> None: + """Close runtimes created by the lightweight test harness.""" + while _TEST_RUNTIMES: + runtime = _TEST_RUNTIMES.pop() + operation = runtime.snapshot().operation + if operation is not None: + runtime.wait_for_terminal(operation.handle, timeout=10.0) + runtime.shutdown(timeout=10.0) diff --git a/dimos/manipulation/control/coordinator_client.py b/dimos/manipulation/control/coordinator_client.py index ed552e0846..6139b7b67c 100644 --- a/dimos/manipulation/control/coordinator_client.py +++ b/dimos/manipulation/control/coordinator_client.py @@ -42,10 +42,12 @@ from __future__ import annotations +import importlib +import importlib.util import math import sys import time -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, cast from dimos.control.components import split_joint_name from dimos.control.coordinator import ControlCoordinator @@ -58,6 +60,12 @@ from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +class _IPythonModule(Protocol): + """Typed subset of the optional IPython module used by this client.""" + + def start_ipython(self, *, argv: list[str], user_ns: dict[str, Any]) -> None: ... + + class CoordinatorClient: """ RPC client for the ControlCoordinator. @@ -560,7 +568,9 @@ def accel(self, value: float | None = None) -> None: def interactive_mode(client: CoordinatorClient, initial_task: str) -> None: """Start IPython interactive mode.""" - import IPython + if importlib.util.find_spec("IPython") is None: + raise RuntimeError("IPython is required for interactive mode") + ipython = cast("_IPythonModule", importlib.import_module("IPython")) shell = CoordinatorShell(client, initial_task) @@ -571,7 +581,7 @@ def interactive_mode(client: CoordinatorClient, initial_task: str) -> None: print("\nType help() for available commands") print("=" * 60 + "\n") - IPython.start_ipython( # type: ignore[no-untyped-call] + ipython.start_ipython( argv=[], user_ns={ "help": shell.help, diff --git a/dimos/manipulation/execution_runtime.py b/dimos/manipulation/execution_runtime.py new file mode 100644 index 0000000000..33d35719cd --- /dev/null +++ b/dimos/manipulation/execution_runtime.py @@ -0,0 +1,2084 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Owner-thread manipulation execution runtime. + +Coordinator calls are effects. They never mutate runtime state; they only +post correlated events back to the owner. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field, replace +from enum import Enum +import queue +import threading +import time +from typing import Any, Generic, Protocol, TypeVar, cast +import uuid + +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +class LifecycleState(str, Enum): + IDLE = "IDLE" + PLANNING = "PLANNING" + READY = "READY" + DISPATCHING = "DISPATCHING" + RUNNING = "RUNNING" + CANCELLING = "CANCELLING" + FAULT = "FAULT" + + +class TaskActivity(str, Enum): + NOT_STARTED = "not_started" + EXECUTE_UNRESOLVED = "execute_unresolved" + ACTIVE = "active" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REMOTE_FAULT = "remote_fault" + UNKNOWN = "unknown" + + +class ActionMethod(str, Enum): + EXECUTE = "execute" + CANCEL = "cancel" + STATUS = "get_state" + RESET = "reset" + GRIPPER_SET = "set_gripper_position" + GRIPPER_GET = "get_gripper_position" + + +class Outcome(str, Enum): + ACCEPTED = "accepted" + RUNNING = "running" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REJECTED = "rejected" + FAILED = "failed" + UNKNOWN = "unknown" + + +class ShutdownState(str, Enum): + OPEN = "open" + CLOSING = "closing" + CLOSED = "closed" + + +class _GatewayGateState(str, Enum): + OPEN = "open" + CANCEL_ONLY = "cancel_only" + SEALED = "sealed" + STOPPED = "stopped" + + +@dataclass(frozen=True) +class ExecutionTopology: + """Canonical group -> robot -> coordinator-task routing.""" + + routes: tuple[tuple[str, str, str], ...] + robot_configs: tuple[RobotModelConfig, ...] = () + inverse_joint_mappings: tuple[tuple[str, tuple[tuple[str, str], ...]], ...] = () + + @classmethod + def from_robot_configs( + cls, + robot_configs: Iterable[RobotModelConfig], + group_routes: Mapping[str, tuple[str, str]] | Iterable[tuple[str, str, str]] | None = None, + group_resolver: Callable[[str], tuple[str, str]] | None = None, + ) -> ExecutionTopology: + """Build canonical routes while retaining the supplied config references.""" + configs = tuple(robot_configs) + by_robot: dict[str, RobotModelConfig] = {} + for config in configs: + if config.name in by_robot: + raise ValueError(f"duplicate robot config: {config.name}") + by_robot[config.name] = config + explicit: dict[str, tuple[str, str]] = {} + if isinstance(group_routes, Mapping): + for group, route in group_routes.items(): + if not isinstance(group, str) or not isinstance(route, tuple) or len(route) != 2: + raise ValueError("invalid group route mapping") + explicit[group] = (str(route[0]), str(route[1])) + elif group_routes is not None: + for group, robot, task in group_routes: + if group in explicit: + raise ValueError(f"duplicate group route: {group}") + explicit[group] = (robot, task) + routes: list[tuple[str, str, str]] = [] + inverse: list[tuple[str, tuple[tuple[str, str], ...]]] = [] + for config in configs: + local_names = set(config.joint_names) + local_to_coordinator: dict[str, str] = {} + for coordinator_name, local_name in config.joint_name_mapping.items(): + if local_name not in local_names: + raise ValueError( + f"mapped local joint is not configured: {config.name}/{local_name}" + ) + previous = local_to_coordinator.get(local_name) + if previous is not None and previous != coordinator_name: + raise ValueError( + f"duplicate local joint mapping for {config.name}/{local_name}" + ) + local_to_coordinator[local_name] = coordinator_name + coordinator_names: set[str] = set() + for local_name in config.joint_names: + coordinator_name = local_to_coordinator.get(local_name, local_name) + if coordinator_name in coordinator_names: + raise ValueError( + f"duplicate coordinator joint mapping for {config.name}/{coordinator_name}" + ) + coordinator_names.add(coordinator_name) + inverse.append( + ( + config.name, + tuple( + (local, local_to_coordinator.get(local, local)) + for local in config.joint_names + ), + ) + ) + for definition in config.planning_groups: + if not set(definition.joint_names).issubset(local_names): + raise ValueError( + f"planning group joints are not configured: {config.name}/{definition.name}" + ) + group_id = f"{config.name}/{definition.name}" + if group_resolver is not None: + route_robot, route_task = group_resolver(group_id) + else: + if not config.coordinator_task_name and group_id not in explicit: + continue + route_robot, route_task = explicit.get( + group_id, (config.name, config.coordinator_task_name or "") + ) + if route_robot != config.name or route_task != config.coordinator_task_name: + raise ValueError(f"route/config mismatch for {group_id}") + routes.append((group_id, route_robot, route_task)) + if explicit and set(explicit) != {group for group, _, _ in routes}: + raise ValueError("route set does not match configured planning groups") + return cls(tuple(routes), configs, tuple(inverse)) + + from_configs = from_robot_configs + from_robot_model_configs = from_robot_configs + + def config_for_robot(self, robot_name: str) -> RobotModelConfig: + for config in self.robot_configs: + if config.name == robot_name: + return config + raise KeyError(f"unknown robot: {robot_name}") + + def coordinator_joint_name(self, robot_name: str, local_name: str) -> str: + for name, mappings in self.inverse_joint_mappings: + if name == robot_name: + for local, coordinator in mappings: + if local == local_name: + return coordinator + break + return local_name + + def route_for_group(self, group_id: str) -> tuple[str, str]: + routes = {group: (robot, task) for group, robot, task in self.routes} + try: + return routes[group_id] + except KeyError as exc: + raise ValueError(f"missing route for planning group {group_id}") from exc + + def affected_robots(self, group_ids: Iterable[str]) -> tuple[str, ...]: + robots: list[str] = [] + for group_id in group_ids: + robot_name, _task_name = self.route_for_group(group_id) + if robot_name in robots: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + robots.append(robot_name) + return tuple(robots) + + @property + def groups(self) -> tuple[str, ...]: + return tuple(group for group, _, _ in self.routes) + + +@dataclass(frozen=True) +class TaskEntry: + planning_group: str + task_name: str + request: Any + robot_name: str | None = None + + +@dataclass(frozen=True) +class ExecutionPlan: + generated_plan: Any + entries: tuple[TaskEntry, ...] + + +@dataclass(frozen=True) +class PreparedPlan: + generated_plan: GeneratedPlan + entries: tuple[TaskEntry, ...] + topology: ExecutionTopology | None = None + _canonical: bool = field(default=False, repr=False, compare=False) + + +def prepare_generated_plan( + generated_plan: GeneratedPlan, topology: ExecutionTopology +) -> PreparedPlan: + """Purely bind a generated multi-robot plan to canonical coordinator tasks.""" + if not isinstance(generated_plan, GeneratedPlan): + raise TypeError("expected GeneratedPlan") + group_ids = tuple(generated_plan.group_ids) + if not group_ids or len(set(group_ids)) != len(group_ids): + raise ValueError("generated plan must contain unique planning groups") + route_by_group = {group: (robot, task) for group, robot, task in topology.routes} + if len(route_by_group) != len(topology.routes): + raise ValueError("topology contains duplicate group routes") + selected_routes: list[tuple[str, str, str, RobotModelConfig, tuple[str, ...]]] = [] + affected: set[str] = set() + for group_id in group_ids: + route = route_by_group.get(group_id) + if route is None: + robot_name = group_id.split("/", 1)[0] + try: + if not topology.config_for_robot(robot_name).coordinator_task_name: + raise ValueError(f"no coordinator task for selected robot {robot_name}") + except KeyError: + pass + raise ValueError(f"missing route for planning group {group_id}") + robot_name, task_name = route + if not task_name: + raise ValueError(f"missing coordinator task for {group_id}") + if robot_name in affected: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + config = topology.config_for_robot(robot_name) + prefix, _, definition_name = group_id.partition("/") + if prefix != robot_name or not definition_name: + raise ValueError(f"invalid canonical planning group {group_id}") + definitions = [ + definition + for definition in config.planning_groups + if definition.name == definition_name + ] + if len(definitions) != 1: + raise ValueError(f"missing config group for {group_id}") + expected_local = tuple(definitions[0].joint_names) + if len(set(expected_local)) != len(expected_local): + raise ValueError(f"duplicate intended joints for {group_id}") + affected.add(robot_name) + selected_routes.append((group_id, robot_name, task_name, config, expected_local)) + expected_by_robot = {robot: set(expected) for _, robot, _, _, expected in selected_routes} + + trajectory = generated_plan.trajectory + names = tuple(trajectory.joint_names) + if not names: + raise ValueError("generated trajectory has no joint names") + if len(set(names)) != len(names): + raise ValueError("duplicate global trajectory joint assignment") + columns: dict[str, dict[str, int]] = {robot: {} for robot in affected} + known_robots = {config.name for config in topology.robot_configs} + for index, global_name in enumerate(names): + if global_name.count("/") != 1: + raise ValueError(f"malformed global joint name: {global_name}") + robot_name, local_name = global_name.split("/", 1) + if not robot_name or not local_name or robot_name not in known_robots: + raise ValueError(f"unknown global joint assignment: {global_name}") + if robot_name in columns: + if local_name not in expected_by_robot[robot_name]: + raise ValueError(f"additional joint column for selected robot: {global_name}") + if local_name in columns[robot_name]: + raise ValueError(f"duplicate robot joint assignment: {global_name}") + columns[robot_name][local_name] = index + points = tuple(trajectory.points) + for point in points: + if len(point.positions) != len(names): + raise ValueError("trajectory point dimension does not match joint names") + if point.velocities and len(point.velocities) != len(names): + raise ValueError("trajectory velocity dimension does not match joint names") + + entries: list[TaskEntry] = [] + for group_id, robot_name, task_name, _config, expected_local in selected_routes: + robot_columns = columns[robot_name] + missing = [local for local in expected_local if local not in robot_columns] + if missing: + raise ValueError(f"missing robot joints for {robot_name}: {missing}") + intended = set(expected_local) + local_order = tuple(local for local in robot_columns if local in intended) + coordinator_names = tuple( + topology.coordinator_joint_name(robot_name, local) for local in local_order + ) + if len(set(coordinator_names)) != len(coordinator_names): + raise ValueError(f"duplicate coordinator joint names for {robot_name}") + local_trajectory = JointTrajectory( + joint_names=list(coordinator_names), + points=[ + TrajectoryPoint( + time_from_start=point.time_from_start, + positions=[point.positions[robot_columns[local]] for local in local_order], + velocities=( + [point.velocities[robot_columns[local]] for local in local_order] + if point.velocities + else None + ), + ) + for point in points + ], + ) + entries.append(TaskEntry(group_id, task_name, {"trajectory": local_trajectory}, robot_name)) + return PreparedPlan(generated_plan, tuple(entries), topology, True) + + +prepare_plan = prepare_generated_plan +prepare_execution_plan = prepare_generated_plan +materialize_prepared_plan = prepare_generated_plan + +PlanInput = ExecutionPlan | PreparedPlan + + +@dataclass(frozen=True) +class OperationHandle: + plan_id: str + operation_id: str + attempt_id: str + + +ExecutionHandle = OperationHandle + + +@dataclass(frozen=True) +class ActionRecord: + action_id: str + method: ActionMethod + started: float + deadline: float + deadline_reported: bool = False + reset_id: str | None = None + + +@dataclass(frozen=True) +class TaskRecord: + task_id: str + task_name: str + entry: TaskEntry + activity: TaskActivity + action: ActionRecord | None = None + cancel_required: bool = False + reset_required: bool = False + + +@dataclass(frozen=True) +class Operation: + handle: OperationHandle + plan: PlanInput + tasks: tuple[TaskRecord, ...] + next_index: int = 0 + cancel_requested: bool = False + cleanup_required: bool = False + uncertain: bool = False + rejected: bool = False + failed: bool = False + diagnostic: str = "" + + +@dataclass(frozen=True) +class ExecutionResult: + handle: OperationHandle + outcome: Outcome + diagnostic: str = "" + + +@dataclass(frozen=True) +class ResetHandle: + reset_id: str + + +@dataclass(frozen=True) +class ResetResult: + handle: ResetHandle + success: bool + diagnostic: str = "" + + +@dataclass(frozen=True) +class ShutdownResult: + success: bool + diagnostic: str = "" + + +@dataclass(frozen=True) +class RuntimeContext: + state: LifecycleState = LifecycleState.IDLE + ready_plan: PlanInput | None = None + ready_plan_id: str | None = None + planning_token: str | None = None + active: Operation | None = None + fault: str | None = None + diagnostic: str | None = None + shutdown: ShutdownState = ShutdownState.OPEN + reset_handle: ResetHandle | None = None + reset_result: ResetResult | None = None + retired_plan_ids: frozenset[str] = frozenset() + dispatch_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + terminal_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + reset_results: tuple[tuple[ResetHandle, ResetResult], ...] = () + reset_proven_inactive: frozenset[str] = frozenset() + reset_completed_tasks: frozenset[str] = frozenset() + reset_deadline: float | None = None + shutdown_result: ShutdownResult | None = None + shutdown_deadline: float | None = None + revision: int = 0 + + +@dataclass(frozen=True) +class RuntimeSnapshot: + state: LifecycleState + ready_plan: PlanInput | None + ready_plan_id: str | None + planning_token: str | None + operation: Operation | None + fault: str | None + diagnostic: str | None + shutdown: ShutdownState + shutdown_result: ShutdownResult | None + revision: int + + +T = TypeVar("T") + + +@dataclass(frozen=True) +class CommandResult(Generic[T]): + accepted: bool + value: T | None = None + diagnostic: str = "" + snapshot: RuntimeSnapshot | None = None + + +class CoordinatorGateway(Protocol): + def execute(self, task_name: str, request: Any) -> Outcome: ... + def cancel(self, task_name: str) -> Outcome: ... + def status(self, task_name: str) -> Outcome: ... + def reset(self, task_name: str) -> Outcome: ... + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: ... + def get_gripper_position(self, hardware_id: str) -> float | None: ... + def stop(self) -> None: ... + + +class ControlCoordinatorGateway: + """Adapter for ControlCoordinator.task_invoke(task, method, kwargs).""" + + def __init__(self, client: Any) -> None: + self._client = client + + def _invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + return self._client.task_invoke(task, method, kwargs) + + def execute(self, task_name: str, request: Any) -> Outcome: + try: + value = self._invoke( + task_name, + "execute", + request if isinstance(request, dict) else {"trajectory": request}, + ) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def cancel(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "cancel", {}) + return ( + Outcome.CANCELLED + if value is True + else Outcome.INACTIVE + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def status(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "get_state", {}) + value = getattr(value, "state", value) + state = TrajectoryState(value) + except (TypeError, ValueError, KeyError): + return Outcome.UNKNOWN + except Exception: + return Outcome.UNKNOWN + return { + TrajectoryState.IDLE: Outcome.INACTIVE, + TrajectoryState.EXECUTING: Outcome.RUNNING, + TrajectoryState.COMPLETED: Outcome.COMPLETED, + TrajectoryState.ABORTED: Outcome.CANCELLED, + TrajectoryState.FAULT: Outcome.FAILED, + }[state] + + def reset(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "reset", {}) + return ( + Outcome.INACTIVE + if value is True + else Outcome.RUNNING + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + try: + value = self._client.set_gripper_position(hardware_id, position) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def get_gripper_position(self, hardware_id: str) -> float | None: + try: + value = self._client.get_gripper_position(hardware_id) + return float(value) if isinstance(value, (int, float)) else None + except Exception: + return None + + def stop(self) -> None: + stop = getattr(self._client, "stop_rpc_client", None) + if callable(stop): + stop() + + +@dataclass(frozen=True) +class _Event: + kind: str + operation_id: str | None = None + attempt_id: str | None = None + task_id: str | None = None + action_id: str | None = None + method: ActionMethod | None = None + outcome: Outcome | None = None + diagnostic: str = "" + payload: Any = None + reset_id: str | None = None + + +@dataclass +class _Call: + function: Callable[[], Any] + answer: queue.Queue[Any] = field(default_factory=lambda: queue.Queue(1)) + + +class ExecutionRuntime: + """All context commits and effect launches happen on one owner thread.""" + + def __init__( + self, + gateway_factory: Callable[[], CoordinatorGateway], + *, + topology: ExecutionTopology | None = None, + plan_validator: Callable[[PlanInput], bool] | None = None, + action_timeout: float = 1.0, + poll_interval: float = 0.1, + ) -> None: + self._gateway = gateway_factory() + self._topology = topology + self._validate_plan = plan_validator or self._valid_plan + self._timeout = action_timeout + self._poll_interval = poll_interval + self._context = RuntimeContext() + self._effects: list[tuple[Operation, TaskRecord, ActionMethod]] = [] + self._context_lock = threading.Lock() + self._events: queue.Queue[_Call | None] = queue.Queue() + self._closed = False + self._stop = threading.Event() + self._condition = threading.Condition(self._context_lock) + self._owner = threading.Thread( + target=self._loop, daemon=True, name="manipulation-execution-owner" + ) + self._owner.start() + self._poller = threading.Thread( + target=self._poll_loop, daemon=True, name="manipulation-execution-poller" + ) + self._poller.start() + self._shutdown_lock = threading.Lock() + self._auxiliary_lock = threading.Lock() + self._auxiliary: dict[threading.Thread, threading.Event] = {} + self._closing_requested = False + self._gateway_condition = threading.Condition() + self._gateway_gate = _GatewayGateState.OPEN + self._gateway_close_started = False + self._gateway_close_done = False + self._gateway_close_error: str | None = None + self._rpc_inflight = 0 + self._admitted_actions: set[str] = set() + self._started_actions: set[str] = set() + + def _start_auxiliary( + self, + target: Callable[..., None], + *args: Any, + name: str, + ) -> None: + """Start a tracked effect thread with cooperative cancellation.""" + cancel = threading.Event() + + def run() -> None: + try: + target(cancel, *args) + finally: + with self._auxiliary_lock: + self._auxiliary.pop(threading.current_thread(), None) + + thread = threading.Thread(target=run, name=name, daemon=True) + with self._auxiliary_lock: + self._auxiliary[thread] = cancel + thread.start() + + def _stop_auxiliary(self, deadline: float) -> bool: + """Cancel and join tracked effects until the caller's deadline.""" + with self._auxiliary_lock: + pending = tuple(self._auxiliary.items()) + for _thread, cancel in pending: + cancel.set() + for thread, _cancel in pending: + thread.join(max(0.0, deadline - time.monotonic())) + with self._auxiliary_lock: + return not self._auxiliary + + def _valid_plan(self, plan: PlanInput) -> bool: + if not isinstance(plan, (ExecutionPlan, PreparedPlan)) or not plan.entries: + return False + if isinstance(plan, ExecutionPlan) and isinstance(plan.generated_plan, GeneratedPlan): + return False + groups = tuple(getattr(plan.generated_plan, "group_ids", ())) + represented = tuple(entry.planning_group for entry in plan.entries) + if len(set(represented)) != len(represented) or len( + set(entry.task_name for entry in plan.entries) + ) != len(plan.entries): + return False + if isinstance(plan, PreparedPlan) and ( + not plan._canonical + or not isinstance(plan.generated_plan, GeneratedPlan) + or plan.topology is None + ): + return False + if isinstance(plan, PreparedPlan) and groups != represented: + return False + if not isinstance(plan, PreparedPlan) and groups and set(groups) != set(represented): + return False + topology = plan.topology if isinstance(plan, PreparedPlan) else self._topology + if topology is not None: + routes = {group: (robot, task) for group, robot, task in topology.routes} + selected_groups = groups or represented + if len(routes) != len(topology.routes) or any( + group not in routes for group in selected_groups + ): + return False + selected = tuple(routes[group] for group in selected_groups) + if tuple(entry.planning_group for entry in plan.entries) != selected_groups: + return False + if any( + ( + route[1] != entry.task_name + or (entry.robot_name is not None and route[0] != entry.robot_name) + ) + for entry, route in zip(plan.entries, selected, strict=True) + ): + return False + return True + + def snapshot(self) -> RuntimeSnapshot: + with self._context_lock: + c = self._context + return RuntimeSnapshot( + c.state, + c.ready_plan, + c.ready_plan_id, + c.planning_token, + c.active, + c.fault, + c.diagnostic, + c.shutdown, + c.shutdown_result, + c.revision, + ) + + def _loop(self) -> None: + while True: + call = self._events.get() + if call is None: + return + try: + value = call.function() + self._process_effects() + call.answer.put(value) + except BaseException as exc: + call.answer.put(exc) + + def _submit(self, fn: Callable[[], T], *, allow_closing: bool = False) -> T: + with self._context_lock: + reject = self._closed or ( + not allow_closing + and (self._closing_requested or self._context.shutdown != ShutdownState.OPEN) + ) + if reject: + diagnostic = "runtime is closed" if self._closed else "runtime is closing" + return cast( + "T", + CommandResult(False, diagnostic=diagnostic, snapshot=self._snapshot_unlocked()), + ) + call = _Call(fn) + self._events.put(call) + result = call.answer.get() + if isinstance(result, BaseException): + raise result + return cast("T", result) + + def _snapshot_unlocked(self) -> RuntimeSnapshot: + c = self._context + return RuntimeSnapshot( + c.state, + c.ready_plan, + c.ready_plan_id, + c.planning_token, + c.active, + c.fault, + c.diagnostic, + c.shutdown, + c.shutdown_result, + c.revision, + ) + + def _commit(self, context: RuntimeContext) -> RuntimeContext: + with self._condition: + self._context = replace(context, revision=context.revision + 1) + self._condition.notify_all() + return self._context + + def _record_dispatch_result_if_absent(self, result: ExecutionResult) -> None: + """Record the first dispatch outcome; later cleanup never replaces it.""" + if dict(self._context.dispatch_results).get(result.handle) is None: + self._commit( + replace( + self._context, + dispatch_results=(*self._context.dispatch_results, (result.handle, result))[ + -16: + ], + ) + ) + + def _reset_pending(self) -> bool: + handle = self._context.reset_handle + return ( + self._context.state == LifecycleState.FAULT + and handle is not None + and handle not in dict(self._context.reset_results) + ) + + def _finish_reset(self, success: bool, diagnostic: str = "") -> None: + handle = self._context.reset_handle + if handle is None or handle in dict(self._context.reset_results): + return + result = ResetResult(handle, success, diagnostic) + reset_results = (*self._context.reset_results, (handle, result))[-16:] + self._commit( + replace( + self._context, + state=LifecycleState.IDLE if success else LifecycleState.FAULT, + active=None if success else self._context.active, + fault=None if success else self._context.fault, + diagnostic=None if success else diagnostic, + reset_results=reset_results, + # A failed attempt is retained in reset_results for its + # caller, but must not remain the active reconciliation gate: + # a later reset command must be able to start a fresh attempt. + reset_handle=None, + reset_proven_inactive=frozenset(), + reset_completed_tasks=frozenset(), + reset_deadline=None, + ) + ) + + def _finish_shutdown(self, success: bool, diagnostic: str = "") -> None: + if self._context.shutdown_result is not None: + return + self._commit(replace(self._context, shutdown_result=ShutdownResult(success, diagnostic))) + + def _maybe_finish_shutdown(self) -> None: + if ( + self._context.shutdown != ShutdownState.CLOSING + or self._context.shutdown_result is not None + ): + return + operation = self._context.active + if operation is None: + self._finish_shutdown(True) + return + if any(task.action is not None for task in operation.tasks): + return + if any(task.activity == TaskActivity.REMOTE_FAULT for task in operation.tasks): + self._finish_shutdown(False, "shutdown safety unresolved: remote fault") + return + pending = next( + ( + task + for task in operation.tasks + if task.activity + in ( + TaskActivity.NOT_STARTED, + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + ) + ), + None, + ) + if pending is not None: + self._schedule_action(operation, pending, ActionMethod.CANCEL) + return + if all( + task.activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED) + for task in operation.tasks + ): + self._finish_shutdown(True) + + def _reconcile_reset(self, op: Operation | None, event: _Event | None = None) -> None: + if not self._reset_pending(): + return + op = op or self._context.active + if op is None: + self._finish_reset(False, "fault has no operation to reconcile") + return + if ( + self._context.reset_deadline is not None + and time.monotonic() >= self._context.reset_deadline + ): + self._finish_reset(False, "fault reset deadline exceeded") + return + if any(task.action is not None for task in op.tasks): + return + + proven = self._context.reset_proven_inactive + if event is not None and event.method == ActionMethod.STATUS: + if event.outcome in (Outcome.INACTIVE, Outcome.CANCELLED, Outcome.COMPLETED): + proven = proven | {event.task_id} if event.task_id is not None else proven + elif event.outcome in (Outcome.RUNNING, Outcome.ACCEPTED): + task = next(task for task in op.tasks if task.task_id == event.task_id) + self._schedule_action(self._context.active or op, task, ActionMethod.CANCEL) + return + else: + self._finish_reset(False, "status reconciliation is unsafe") + return + elif event is not None and event.method == ActionMethod.RESET: + if event.outcome != Outcome.INACTIVE: + self._finish_reset(False, "remote reset failed") + return + if event.task_id is not None: + self._commit( + replace( + self._context, + reset_completed_tasks=self._context.reset_completed_tasks | {event.task_id}, + ) + ) + elif event is not None and event.outcome in (Outcome.UNKNOWN, Outcome.FAILED): + self._finish_reset(False, "fault reconciliation is unsafe") + return + + if proven != self._context.reset_proven_inactive: + self._commit(replace(self._context, reset_proven_inactive=proven)) + + current = self._context.active or op + unproven = next( + ( + task + for task in current.tasks + if task.task_id not in self._context.reset_proven_inactive + ), + None, + ) + if unproven is not None: + method = ( + ActionMethod.CANCEL + if unproven.activity + in ( + TaskActivity.NOT_STARTED, + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + ) + else ActionMethod.STATUS + ) + self._schedule_action(current, unproven, method) + return + + pending_reset = next( + ( + task + for task in current.tasks + if task.reset_required and task.task_id not in self._context.reset_completed_tasks + ), + None, + ) + if pending_reset is not None: + self._schedule_action(current, pending_reset, ActionMethod.RESET) + return + self._finish_reset(True) + + def _reduce(self, event: _Event) -> Any: + """Pure transition function plus immutable context commit.""" + c = self._context + if c.shutdown == ShutdownState.CLOSED: + return False + if c.shutdown == ShutdownState.CLOSING and event.kind not in ( + "action_deadline", + "action_finished", + "effect_rejected", + "shutdown", + "shutdown_deadline", + ): + return False + if event.kind == "start_planning" and c.state in ( + LifecycleState.IDLE, + LifecycleState.READY, + ): + token = str(uuid.uuid4()) + self._commit( + replace( + c, + state=LifecycleState.PLANNING, + planning_token=token, + ready_plan=None, + ready_plan_id=None, + diagnostic=None, + ) + ) + return token + if event.kind == "cancel_planning" and c.state == LifecycleState.PLANNING: + self._commit(replace(c, state=LifecycleState.IDLE, planning_token=None)) + return True + if ( + event.kind == "plan_ready" + and c.state == LifecycleState.PLANNING + and event.payload is not None + ): + plan_id = str(uuid.uuid4()) + self._commit( + replace( + c, + state=LifecycleState.READY, + ready_plan=event.payload, + ready_plan_id=plan_id, + planning_token=None, + diagnostic=None, + ) + ) + return plan_id + if event.kind == "planning_failed" and c.state == LifecycleState.PLANNING: + self._commit( + replace( + c, + state=LifecycleState.IDLE, + planning_token=None, + ready_plan=None, + ready_plan_id=None, + diagnostic=event.diagnostic or "planning failed", + ) + ) + return True + if event.kind == "clear_ready_plan" and c.state == LifecycleState.READY: + self._commit(replace(c, state=LifecycleState.IDLE, ready_plan=None, ready_plan_id=None)) + return True + if event.kind == "dispatch_start" and c.state in ( + LifecycleState.IDLE, + LifecycleState.READY, + ): + handle, plan = event.payload + self._commit( + replace( + c, + state=LifecycleState.DISPATCHING, + ready_plan=None, + ready_plan_id=None, + active=Operation( + handle, + plan, + tuple( + TaskRecord(str(uuid.uuid4()), e.task_name, e, TaskActivity.NOT_STARTED) + for e in plan.entries + ), + ), + fault=None, + diagnostic=None, + reset_handle=None, + reset_proven_inactive=frozenset(), + reset_completed_tasks=frozenset(), + reset_deadline=None, + shutdown_result=None, + shutdown_deadline=None, + ) + ) + return True + if event.kind == "reset_deadline": + reset_handle = self._context.reset_handle + if ( + self._reset_pending() + and reset_handle is not None + and event.payload == reset_handle.reset_id + ): + self._finish_reset(False, "fault reset deadline exceeded") + return True + return False + if event.kind == "shutdown_deadline": + if ( + c.shutdown == ShutdownState.CLOSING + and c.shutdown_result is None + and event.payload == c.shutdown_deadline + ): + self._finish_shutdown(False, "shutdown safety deadline exceeded") + return True + return False + op = c.active + if event.kind == "effect_rejected": + if ( + op is None + or event.operation_id != op.handle.operation_id + or event.attempt_id != op.handle.attempt_id + ): + return False + task_index = next( + (i for i, task in enumerate(op.tasks) if task.task_id == event.task_id), None + ) + if task_index is None: + return False + task = op.tasks[task_index] + if ( + task.action is None + or task.action.action_id != event.action_id + or task.action.method != event.method + ): + return False + if event.reset_id is not None: + current_reset_id = c.reset_handle.reset_id if c.reset_handle is not None else None + if event.reset_id != current_reset_id: + tasks = list(op.tasks) + tasks[task_index] = replace(task, action=None) + self._commit(replace(c, active=replace(op, tasks=tuple(tasks)))) + if self._reset_pending(): + self._reconcile_reset(self._context.active) + return True + activity = ( + TaskActivity.INACTIVE if event.method == ActionMethod.EXECUTE else task.activity + ) + tasks = list(op.tasks) + tasks[task_index] = replace(task, action=None, activity=activity) + self._commit(replace(c, active=replace(op, tasks=tuple(tasks)))) + if c.shutdown == ShutdownState.CLOSING and c.shutdown_result is None: + self._finish_shutdown(False, "shutdown dropped an unresolved coordinator effect") + return True + if event.kind in ("action_deadline", "action_finished"): + if ( + op is None + or event.operation_id != op.handle.operation_id + or event.attempt_id != op.handle.attempt_id + or event.task_id not in {t.task_id for t in op.tasks} + ): + return False + task_index = next(i for i, task in enumerate(op.tasks) if task.task_id == event.task_id) + task = op.tasks[task_index] + if ( + task.action is None + or task.action.action_id != event.action_id + or task.action.method != event.method + ): + return False + if event.reset_id is not None: + current_reset_id = c.reset_handle.reset_id if c.reset_handle is not None else None + if event.reset_id != current_reset_id: + if event.kind == "action_deadline": + return False + tasks = list(op.tasks) + tasks[task_index] = replace(task, action=None) + refreshed = replace(op, tasks=tuple(tasks)) + self._commit(replace(c, active=refreshed)) + self._reconcile_reset(refreshed) + return True + event_detail = f"task={task.task_name} action={event.action_id} outcome={event.outcome.value if event.outcome is not None else 'unknown'}" + if event.kind == "action_deadline": + if task.action.deadline_reported: + return False + if not self._action_was_started(task.action.action_id): + return self._reduce(replace(event, kind="effect_rejected")) + action = replace(task.action, deadline_reported=True) + tasks = list(op.tasks) + tasks[task_index] = replace( + task, action=action, activity=TaskActivity.UNKNOWN, cancel_required=True + ) + with self._gateway_condition: + self._started_actions.discard(task.action.action_id) + fault_op = replace(op, tasks=tuple(tasks), uncertain=True) + result = ExecutionResult(op.handle, Outcome.UNKNOWN, "coordinator action deadline") + self._commit( + replace( + c, + state=LifecycleState.FAULT, + fault="coordinator action deadline", + diagnostic="coordinator action deadline", + active=fault_op, + ) + ) + if c.shutdown == ShutdownState.CLOSING: + self._finish_shutdown(False, "shutdown safety unresolved: action deadline") + return True + if self._reset_pending(): + self._finish_reset(False, "fault reconciliation action deadline") + return True + self._record_dispatch_result_if_absent(result) + self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) + return True + tasks = list(op.tasks) + with self._gateway_condition: + self._started_actions.discard(task.action.action_id) + if event.outcome in (Outcome.ACCEPTED, Outcome.RUNNING): + activity = TaskActivity.ACTIVE + elif event.outcome == Outcome.COMPLETED: + activity = TaskActivity.COMPLETED + elif event.outcome == Outcome.CANCELLED: + activity = ( + TaskActivity.CANCELLED + if event.method == ActionMethod.CANCEL or c.state == LifecycleState.CANCELLING + else TaskActivity.UNKNOWN + ) + elif event.outcome == Outcome.INACTIVE: + activity = TaskActivity.INACTIVE + elif event.outcome == Outcome.FAILED: + activity = TaskActivity.REMOTE_FAULT + elif event.outcome == Outcome.REJECTED and event.method == ActionMethod.EXECUTE: + activity = TaskActivity.INACTIVE + else: + activity = TaskActivity.UNKNOWN + tasks[task_index] = replace( + task, + activity=activity, + action=None, + cancel_required=task.cancel_required + or (c.state in (LifecycleState.CANCELLING, LifecycleState.FAULT)), + reset_required=task.reset_required or event.outcome == Outcome.FAILED, + ) + next_op = replace(op, tasks=tuple(tasks)) + self._commit(replace(c, active=next_op)) + if c.shutdown == ShutdownState.CLOSING: + if self._context.shutdown_result is None: + self._maybe_finish_shutdown() + return True + if self._reset_pending(): + self._reconcile_reset(next_op, event) + return True + if c.state == LifecycleState.FAULT: + if event.method != ActionMethod.CANCEL: + self._schedule_actions(next_op, next_op.tasks, ActionMethod.CANCEL) + return True + if event.method == ActionMethod.STATUS and event.outcome in ( + Outcome.UNKNOWN, + Outcome.FAILED, + ): + fault_op = replace(next_op, uncertain=True) + diagnostic = f"status is unsafe: {event_detail}" + result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.FAULT, + fault=diagnostic, + diagnostic=diagnostic, + active=fault_op, + ) + ) + self._record_dispatch_result_if_absent(result) + self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) + return True + if event.method == ActionMethod.CANCEL and event.outcome in ( + Outcome.UNKNOWN, + Outcome.FAILED, + ): + fault_op = replace(next_op, uncertain=True) + diagnostic = ( + "cancellation is uncertain" + if event.outcome == Outcome.UNKNOWN + else "cancellation failed" + ) + result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.FAULT, + fault=diagnostic, + diagnostic=diagnostic, + active=fault_op, + ) + ) + self._record_dispatch_result_if_absent(result) + if event.outcome == Outcome.FAILED: + self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) + return True + if ( + c.state == LifecycleState.CANCELLING + and event.method == ActionMethod.STATUS + and event.outcome == Outcome.RUNNING + ): + self._schedule_actions(next_op, (tasks[task_index],), ActionMethod.CANCEL) + return True + if ( + event.method == ActionMethod.STATUS + and event.outcome in (Outcome.INACTIVE, Outcome.CANCELLED) + and c.state == LifecycleState.RUNNING + ): + diagnostic = f"status reported terminal failure: {event_detail}" + failed_op = replace( + next_op, + cancel_requested=True, + cleanup_required=True, + rejected=False, + failed=True, + diagnostic=diagnostic, + ) + self._commit( + replace( + self._context, + state=LifecycleState.CANCELLING, + active=failed_op, + diagnostic=diagnostic, + ) + ) + self._schedule_actions(failed_op, failed_op.tasks, ActionMethod.CANCEL) + return True + if event.method == ActionMethod.EXECUTE and event.outcome in ( + Outcome.REJECTED, + Outcome.UNKNOWN, + Outcome.FAILED, + ): + if event.outcome == Outcome.REJECTED: + skipped = tuple( + replace(item, activity=TaskActivity.INACTIVE) + if item.activity == TaskActivity.NOT_STARTED + else item + for item in next_op.tasks + ) + next_op = replace(next_op, tasks=skipped) + self._commit(replace(self._context, active=next_op)) + if event.outcome != Outcome.REJECTED: + fault_op = replace(next_op, uncertain=True) + diagnostic = f"execute outcome is uncertain: {event_detail}" + result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.FAULT, + fault=diagnostic, + diagnostic=diagnostic, + active=fault_op, + ) + ) + self._record_dispatch_result_if_absent(result) + self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) + elif any( + t.activity in (TaskActivity.ACTIVE, TaskActivity.UNKNOWN) for t in next_op.tasks + ): + diagnostic = f"execute rejected: {event_detail}" + rejected_op = replace( + next_op, + cancel_requested=True, + rejected=True, + cleanup_required=True, + diagnostic=diagnostic, + ) + rejection_result = ExecutionResult(op.handle, Outcome.REJECTED, diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.CANCELLING, + active=rejected_op, + diagnostic=diagnostic, + ) + ) + self._record_dispatch_result_if_absent(rejection_result) + self._schedule_actions(rejected_op, rejected_op.tasks, ActionMethod.CANCEL) + else: + diagnostic = f"execute rejected: {event_detail}" + result = ExecutionResult(next_op.handle, Outcome.REJECTED, diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.IDLE, + active=None, + diagnostic=diagnostic, + terminal_results=( + *self._context.terminal_results, + (next_op.handle, result), + )[-16:], + ) + ) + self._record_dispatch_result_if_absent(result) + return True + self._advance(next_op, event) + return True + if ( + ( + event.kind == "cancel" + or ( + event.kind == "cancel_if_current" + and op is not None + and event.payload == op.handle + ) + ) + and op is not None + and c.state in (LifecycleState.DISPATCHING, LifecycleState.RUNNING) + ): + cancel_tasks: tuple[TaskRecord, ...] = tuple( + replace(task, activity=TaskActivity.INACTIVE, cancel_required=True) + if task.activity == TaskActivity.NOT_STARTED + else replace(task, cancel_required=True) + if task.activity == TaskActivity.EXECUTE_UNRESOLVED + else task + for task in op.tasks + ) + next_op = replace(op, tasks=cancel_tasks, cancel_requested=True) + self._commit(replace(c, state=LifecycleState.CANCELLING, active=next_op)) + if c.state == LifecycleState.DISPATCHING and not dict(c.dispatch_results).get( + op.handle + ): + result = ExecutionResult( + op.handle, Outcome.CANCELLED, "cancelled before dispatch completed" + ) + self._record_dispatch_result_if_absent(result) + self._schedule_actions(next_op, cancel_tasks, ActionMethod.CANCEL) + return True + if event.kind == "poll_tick" and op is not None and c.state == LifecycleState.RUNNING: + candidate = next( + ( + item + for item in op.tasks + if item.activity == TaskActivity.ACTIVE and item.action is None + ), + None, + ) + if candidate is not None: + self._schedule_actions(op, (candidate,), ActionMethod.STATUS) + return True + if event.kind == "fault": + diagnostic = event.diagnostic or "remote uncertainty" + self._commit( + replace( + c, + state=LifecycleState.FAULT, + fault=diagnostic, + diagnostic=diagnostic, + active=replace(op, uncertain=True) if op else None, + ) + ) + return True + if event.kind == "shutdown": + if self._reset_pending(): + self._finish_reset(False, "shutdown interrupted fault reset") + c = self._context + op = c.active + if op is not None: + shutdown_tasks: tuple[TaskRecord, ...] = tuple( + replace(task, activity=TaskActivity.INACTIVE, cancel_required=True) + if task.activity == TaskActivity.NOT_STARTED + else replace(task, cancel_required=True) + for task in op.tasks + ) + op = replace(op, tasks=shutdown_tasks, cancel_requested=True, cleanup_required=True) + self._commit( + replace( + c, + shutdown=ShutdownState.CLOSING, + state=LifecycleState.CANCELLING if op else LifecycleState.IDLE, + planning_token=None, + ready_plan=None, + ready_plan_id=None, + active=op, + shutdown_deadline=event.payload, + ) + ) + if op is None: + self._finish_shutdown(True) + else: + self._schedule_actions(op, op.tasks, ActionMethod.CANCEL) + self._maybe_finish_shutdown() + return True + return False + + def _advance(self, op: Operation, event: _Event) -> None: + # This is called only by the owner after a correlated result. + cancelling = op.cancel_requested or self._context.state in ( + LifecycleState.CANCELLING, + LifecycleState.FAULT, + ) + if cancelling: + self._schedule_actions(op, op.tasks, ActionMethod.CANCEL) + if ( + not cancelling + and event.method == ActionMethod.EXECUTE + and event.outcome == Outcome.ACCEPTED + ): + if op.next_index + 1 < len(op.tasks): + self._commit( + replace(self._context, active=replace(op, next_index=op.next_index + 1)) + ) + self._schedule_execute(self._context.active) + return + if all(task.activity == TaskActivity.ACTIVE for task in op.tasks): + result = ExecutionResult(op.handle, Outcome.ACCEPTED) + self._commit(replace(self._context, state=LifecycleState.RUNNING)) + self._record_dispatch_result_if_absent(result) + return + terminal = all( + task.action is None + and task.activity + in (TaskActivity.COMPLETED, TaskActivity.CANCELLED, TaskActivity.INACTIVE) + for task in op.tasks + ) + completed = all( + task.activity == TaskActivity.COMPLETED and task.action is None for task in op.tasks + ) + if terminal and (op.cancel_requested or completed): + if self._context.state == LifecycleState.FAULT: + return + outcome = ( + Outcome.REJECTED + if op.rejected + else Outcome.FAILED + if op.failed + else Outcome.CANCELLED + if op.cancel_requested + else Outcome.COMPLETED + ) + result = ExecutionResult(op.handle, outcome, op.diagnostic) + self._commit( + replace( + self._context, + state=LifecycleState.IDLE, + active=None, + diagnostic=op.diagnostic or self._context.diagnostic, + terminal_results=(*self._context.terminal_results, (op.handle, result))[-16:], + ) + ) + + def _schedule_execute(self, op: Operation | None) -> None: + if op is None or op.next_index >= len(op.tasks): + return + task = op.tasks[op.next_index] + if task.action is not None: + return + self._schedule_action(op, task, ActionMethod.EXECUTE) + + def _schedule_action(self, op: Operation, task: TaskRecord, method: ActionMethod) -> None: + self._schedule_actions(op, (task,), method) + + def _schedule_actions(self, op: Operation, requested: Any, method: ActionMethod) -> None: + """Allocate and commit a whole cleanup batch against one operation.""" + if ( + self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED) + or self._context.shutdown == ShutdownState.CLOSED + or (self._context.shutdown == ShutdownState.CLOSING and method != ActionMethod.CANCEL) + ): + return + candidates = tuple( + task + for task in requested + if task.action is None + and ( + method in (ActionMethod.STATUS, ActionMethod.RESET) + or task.activity + not in ( + TaskActivity.INACTIVE, + TaskActivity.COMPLETED, + TaskActivity.CANCELLED, + TaskActivity.REMOTE_FAULT, + ) + ) + ) + if not candidates: + return + now = time.monotonic() + reset_id = ( + self._context.reset_handle.reset_id + if self._reset_pending() and self._context.reset_handle is not None + else None + ) + actions = { + task.task_id: ActionRecord( + str(uuid.uuid4()), method, now, now + self._timeout, reset_id=reset_id + ) + for task in candidates + } + task_map = {task.task_id: task for task in op.tasks} + for task in candidates: + task_map[task.task_id] = replace( + task, + action=actions[task.task_id], + activity=TaskActivity.EXECUTE_UNRESOLVED + if method == ActionMethod.EXECUTE + else task.activity, + ) + updated = replace(op, tasks=tuple(task_map[task.task_id] for task in op.tasks)) + self._commit(replace(self._context, active=updated)) + self._effects.extend((updated, task_map[task.task_id], method) for task in candidates) + + def _process_effects(self) -> None: + effects, self._effects = self._effects, [] + if self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED): + if ( + self._context.shutdown == ShutdownState.CLOSING + and self._context.shutdown_result is None + and effects + ): + self._finish_shutdown(False, "shutdown gateway closed before pending effects") + return + for op, task, method in effects: + admitted = self._run_action(op, task, method) + if ( + not admitted + and self._context.shutdown == ShutdownState.CLOSING + and self._context.shutdown_result is None + ): + self._finish_shutdown(False, "shutdown dropped an unresolved coordinator effect") + + def _invoke_gateway_action( + self, action_id: str, task_name: str, method: ActionMethod, request: Any + ) -> tuple[bool, Any]: + """Lease a gateway call immediately before invoking it; no context lock crosses I/O.""" + with self._gateway_condition: + allowed = self._gateway_gate == _GatewayGateState.OPEN or ( + self._gateway_gate == _GatewayGateState.CANCEL_ONLY + and method == ActionMethod.CANCEL + ) + if not allowed: + return False, None + self._rpc_inflight += 1 + self._admitted_actions.add(action_id) + self._started_actions.add(action_id) + try: + outcome: Any = Outcome.UNKNOWN + try: + if method == ActionMethod.EXECUTE: + outcome = self._gateway.execute(task_name, request) + elif method == ActionMethod.CANCEL: + outcome = self._gateway.cancel(task_name) + elif method == ActionMethod.STATUS: + outcome = self._gateway.status(task_name) + elif method == ActionMethod.RESET: + outcome = self._gateway.reset(task_name) + elif method == ActionMethod.GRIPPER_SET: + outcome = self._gateway.set_gripper_position(task_name, float(request)) + else: + outcome = self._gateway.get_gripper_position(task_name) + except Exception: + outcome = Outcome.UNKNOWN + return True, outcome + finally: + with self._gateway_condition: + self._rpc_inflight -= 1 + self._admitted_actions.discard(action_id) + self._gateway_condition.notify_all() + + def _action_was_admitted(self, action_id: str) -> bool: + with self._gateway_condition: + return action_id in self._admitted_actions + + def _action_was_started(self, action_id: str) -> bool: + with self._gateway_condition: + return action_id in self._started_actions + + def _run_action(self, op: Operation, task: TaskRecord, method: ActionMethod) -> bool: + action = task.action + if action is None: + return False + request = task.entry.request + + def call(_cancel: threading.Event) -> None: + admitted, outcome = self._invoke_gateway_action( + action.action_id, task.task_name, method, request + ) + if admitted: + self._events.put( + _Call( + lambda: self._reduce( + _Event( + "action_finished", + op.handle.operation_id, + op.handle.attempt_id, + task.task_id, + action.action_id, + method, + outcome, + reset_id=action.reset_id, + ) + ) + ) + ) + else: + self._events.put( + _Call( + lambda: self._reduce( + _Event( + "effect_rejected", + op.handle.operation_id, + op.handle.attempt_id, + task.task_id, + action.action_id, + method, + reset_id=action.reset_id, + ) + ) + ) + ) + + def deadline(cancel: threading.Event) -> None: + if cancel.wait(max(0.0, action.deadline - time.monotonic())): + return + self._events.put( + _Call( + lambda: self._reduce( + _Event( + "action_deadline", + op.handle.operation_id, + op.handle.attempt_id, + task.task_id, + action.action_id, + method, + reset_id=action.reset_id, + ) + ) + ) + ) + + self._start_auxiliary(call, name=f"action-{action.action_id}") + self._start_auxiliary(deadline, name=f"deadline-{action.action_id}") + return True + + def start_planning(self) -> str | None: + return cast("str | None", self._submit(lambda: self._reduce(_Event("start_planning")))) + + def complete_planning(self, token: str, plan: PlanInput) -> CommandResult[str]: + def fn() -> CommandResult[str]: + if self._context.planning_token != token: + return CommandResult( + False, diagnostic="planning token mismatch", snapshot=self.snapshot() + ) + if not self._validate_plan(plan): + self._reduce(_Event("planning_failed", diagnostic="invalid prepared plan")) + return CommandResult( + False, diagnostic="invalid prepared plan", snapshot=self.snapshot() + ) + plan_id = self._reduce(_Event("plan_ready", payload=plan)) + return CommandResult(True, plan_id, snapshot=self.snapshot()) + + return self._submit(fn) + + def fail_planning(self, token: str, diagnostic: str) -> CommandResult[None]: + return self._submit( + lambda: CommandResult( + self._context.planning_token == token, + diagnostic=diagnostic + if self._context.planning_token == token + else "planning token mismatch", + snapshot=self.snapshot(), + ) + if self._context.planning_token != token + else ( + self._reduce(_Event("planning_failed", diagnostic=diagnostic)) + and CommandResult(True, diagnostic=diagnostic, snapshot=self.snapshot()) + ) + ) + + def cancel_planning(self) -> CommandResult[None]: + return self._submit( + lambda: CommandResult( + bool(self._reduce(_Event("cancel_planning"))), snapshot=self.snapshot() + ) + ) + + def clear_ready_plan(self) -> CommandResult[None]: + return self._submit( + lambda: CommandResult( + bool(self._reduce(_Event("clear_ready_plan"))), snapshot=self.snapshot() + ) + ) + + def execute_explicit(self, prepared_plan: PlanInput) -> CommandResult[OperationHandle]: + def fn() -> CommandResult[OperationHandle]: + if self._context.state in ( + LifecycleState.DISPATCHING, + LifecycleState.RUNNING, + LifecycleState.CANCELLING, + ) and self._validate_plan(prepared_plan): + handle = OperationHandle(str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) + result = ExecutionResult(handle, Outcome.REJECTED, "execution already active") + self._commit( + replace( + self._context, + diagnostic="execution already active", + terminal_results=(*self._context.terminal_results, (handle, result))[-16:], + ) + ) + return CommandResult( + False, + handle, + diagnostic="execution already active", + snapshot=self.snapshot(), + ) + if self._context.state not in ( + LifecycleState.IDLE, + LifecycleState.READY, + ) or not self._validate_plan(prepared_plan): + self._commit( + replace(self._context, diagnostic="execution unavailable or invalid plan") + ) + return CommandResult( + False, + diagnostic="execution unavailable or invalid plan", + snapshot=self.snapshot(), + ) + handle = OperationHandle(str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) + if self._reduce(_Event("dispatch_start", payload=(handle, prepared_plan))): + self._schedule_execute(self._context.active) + return CommandResult(True, handle, snapshot=self.snapshot()) + return CommandResult(False, diagnostic="dispatch rejected", snapshot=self.snapshot()) + + return self._submit(fn) + + def execute_ready(self) -> CommandResult[OperationHandle]: + def fn() -> CommandResult[OperationHandle]: + if self._context.ready_plan is None or self._context.ready_plan_id is None: + return CommandResult(False, diagnostic="no ready plan", snapshot=self.snapshot()) + plan = self._context.ready_plan + return self._execute_owner(plan, self._context.ready_plan_id) + + return self._submit(fn) + + def _execute_owner( + self, plan: PlanInput, plan_id: str | None = None + ) -> CommandResult[OperationHandle]: + handle = OperationHandle(plan_id or str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) + if self._reduce(_Event("dispatch_start", payload=(handle, plan))): + self._schedule_execute(self._context.active) + return CommandResult(True, handle, snapshot=self.snapshot()) + return CommandResult(False, diagnostic="dispatch rejected", snapshot=self.snapshot()) + + def cancel(self) -> CommandResult[None]: + return self._submit( + lambda: CommandResult(bool(self._reduce(_Event("cancel"))), snapshot=self.snapshot()) + ) + + def cancel_if_current(self, handle: OperationHandle) -> CommandResult[None]: + """Cancel only if ``handle`` is the current operation at the owner event.""" + + def fn() -> CommandResult[None]: + accepted = bool(self._reduce(_Event("cancel_if_current", payload=handle))) + current = self._context.active + diagnostic = ( + "" + if accepted + else "operation is no longer current" + if current is None or current.handle != handle + else f"cancellation rejected in state {self._context.state.name}" + ) + return CommandResult( + accepted, + diagnostic=diagnostic, + snapshot=self.snapshot(), + ) + + return self._submit(fn) + + def poll(self) -> CommandResult[None]: + return self._submit( + lambda: CommandResult(bool(self._reduce(_Event("poll_tick"))), snapshot=self.snapshot()) + ) + + def set_gripper_position(self, hardware_id: str, position: float) -> CommandResult[Outcome]: + def fn() -> CommandResult[Outcome]: + admitted, outcome = self._invoke_gateway_action( + str(uuid.uuid4()), hardware_id, ActionMethod.GRIPPER_SET, position + ) + if not admitted: + return CommandResult( + False, diagnostic="gateway is closing", snapshot=self.snapshot() + ) + return CommandResult(True, outcome, snapshot=self.snapshot()) + + return self._submit(fn) + + def get_gripper_position(self, hardware_id: str) -> CommandResult[float]: + def fn() -> CommandResult[float]: + admitted, position = self._invoke_gateway_action( + str(uuid.uuid4()), hardware_id, ActionMethod.GRIPPER_GET, None + ) + if not admitted: + return CommandResult( + False, diagnostic="gateway is closing", snapshot=self.snapshot() + ) + if not isinstance(position, (int, float)): + return CommandResult( + False, + diagnostic="gateway returned no gripper position", + snapshot=self.snapshot(), + ) + return CommandResult(True, float(position), snapshot=self.snapshot()) + + return self._submit(fn) + + def wait_for_dispatch( + self, handle: OperationHandle, timeout: float = 1.0 + ) -> CommandResult[ExecutionResult]: + return self._wait(handle, "dispatch", timeout) + + def wait_for_terminal( + self, handle: OperationHandle, timeout: float = 1.0 + ) -> CommandResult[ExecutionResult]: + return self._wait(handle, "terminal", timeout) + + def _wait( + self, handle: OperationHandle, result_kind: str, timeout: float + ) -> CommandResult[ExecutionResult]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + with self._condition: + results = ( + self._context.dispatch_results + if result_kind == "dispatch" + else self._context.terminal_results + ) + found = dict(results).get(handle) + if found is None: + self._condition.wait(min(0.01, deadline - time.monotonic())) + if found is not None: + return CommandResult(True, found, snapshot=self.snapshot()) + return CommandResult(False, diagnostic="timeout", snapshot=self.snapshot()) + + def _poll_loop(self) -> None: + while not self._stop.wait(self._poll_interval): + self._events.put(_Call(lambda: self._reduce(_Event("poll_tick")))) + + def reset(self) -> CommandResult[ResetHandle]: + def fn() -> CommandResult[ResetHandle]: + if self._context.state != LifecycleState.FAULT: + return CommandResult( + False, diagnostic="reset requires FAULT", snapshot=self.snapshot() + ) + handle = self._context.reset_handle or ResetHandle(str(uuid.uuid4())) + if self._context.reset_handle is not None: + return CommandResult(True, self._context.reset_handle, snapshot=self.snapshot()) + deadline = time.monotonic() + max(self._timeout, 0.05) + self._commit( + replace( + self._context, + reset_handle=handle, + reset_deadline=deadline, + reset_proven_inactive=frozenset(), + reset_completed_tasks=frozenset(), + ) + ) + self._reconcile_reset(self._context.active) + self._start_auxiliary( + self._reset_deadline, handle, deadline, name=f"reset-deadline-{handle.reset_id}" + ) + return CommandResult(True, handle, snapshot=self.snapshot()) + + return self._submit(fn) + + def _reset_deadline( + self, cancel: threading.Event, handle: ResetHandle, deadline: float + ) -> None: + if cancel.wait(max(0.0, deadline - time.monotonic())): + return + self._events.put( + _Call( + lambda: self._reduce( + _Event( + "reset_deadline", + payload=handle.reset_id, + reset_id=handle.reset_id, + ) + ) + ) + ) + + def wait_for_reset( + self, handle: ResetHandle, timeout: float = 1.0 + ) -> CommandResult[ResetResult]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + with self._condition: + result = dict(self._context.reset_results).get(handle) + if result is None: + self._condition.wait(min(0.01, deadline - time.monotonic())) + if result is not None: + return CommandResult(True, result, snapshot=self.snapshot()) + return CommandResult(False, diagnostic="timeout", snapshot=self.snapshot()) + + def _seal_gateway(self) -> None: + with self._gateway_condition: + if self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED): + return + self._gateway_gate = _GatewayGateState.SEALED + self._gateway_condition.notify_all() + self._start_auxiliary(self._close_gateway, name="gateway-close") + + def _close_gateway(self, _cancel: threading.Event) -> None: + with self._gateway_condition: + while self._rpc_inflight: + self._gateway_condition.wait() + if self._gateway_close_started: + return + self._gateway_close_started = True + try: + self._gateway.stop() + except BaseException as exc: + with self._gateway_condition: + self._gateway_close_error = f"gateway close failed: {exc}" + self._gateway_condition.notify_all() + return + else: + with self._gateway_condition: + self._gateway_gate = _GatewayGateState.STOPPED + self._gateway_close_done = True + self._gateway_condition.notify_all() + + def _wait_gateway_closed(self, deadline: float) -> bool: + with self._gateway_condition: + while ( + not self._gateway_close_done + and self._gateway_close_error is None + and time.monotonic() < deadline + ): + self._gateway_condition.wait(min(0.01, deadline - time.monotonic())) + return self._gateway_close_done and self._gateway_close_error is None + + def shutdown(self, timeout: float = 1.0) -> CommandResult[ShutdownResult]: + with self._shutdown_lock: + if self._closed: + result = self._context.shutdown_result or ShutdownResult( + False, "runtime is closed without a shutdown result" + ) + return CommandResult(result.success, result, snapshot=self.snapshot()) + if not self._owner.is_alive() and self._context.shutdown == ShutdownState.CLOSING: + result = self._context.shutdown_result or ShutdownResult( + False, "shutdown result unavailable" + ) + if not result.success: + return CommandResult(False, result, snapshot=self.snapshot()) + self._commit( + replace(self._context, shutdown=ShutdownState.CLOSED, shutdown_result=result) + ) + self._closed = True + return CommandResult(result.success, result, snapshot=self.snapshot()) + deadline = time.monotonic() + max(0.0, timeout) + with self._condition: + with self._gateway_condition: + self._closing_requested = True + self._context = replace( + self._context, + shutdown=ShutdownState.CLOSING, + shutdown_deadline=deadline, + revision=self._context.revision + 1, + ) + if self._gateway_gate == _GatewayGateState.OPEN: + self._gateway_gate = _GatewayGateState.CANCEL_ONLY + self._gateway_condition.notify_all() + self._condition.notify_all() + + self._stop.set() + call = _Call(lambda: self._reduce(_Event("shutdown", payload=deadline))) + self._events.put(call) + remaining = max(0.0, deadline - time.monotonic()) + try: + call.answer.get(timeout=remaining) + except queue.Empty: + pass + + with self._condition: + while self._context.shutdown_result is None and time.monotonic() < deadline: + self._condition.wait(min(0.01, deadline - time.monotonic())) + + if self._context.shutdown_result is None: + deadline_call = _Call( + lambda: self._reduce(_Event("shutdown_deadline", payload=deadline)) + ) + self._events.put(deadline_call) + remaining = max(0.0, deadline - time.monotonic()) + if remaining: + try: + deadline_call.answer.get(timeout=remaining) + except queue.Empty: + pass + if self._context.shutdown_result is None: + self._commit( + replace( + self._context, + shutdown_result=ShutdownResult(False, "shutdown safety deadline exceeded"), + ) + ) + + self._seal_gateway() + close_done = self._wait_gateway_closed(deadline) + if not close_done and ( + self._context.shutdown_result is None or self._context.shutdown_result.success + ): + with self._gateway_condition: + close_diagnostic = ( + self._gateway_close_error + or "gateway close unresolved before shutdown deadline" + ) + self._commit( + replace(self._context, shutdown_result=ShutdownResult(False, close_diagnostic)) + ) + auxiliary_drained = self._stop_auxiliary(deadline) + if not auxiliary_drained and ( + self._context.shutdown_result is None or self._context.shutdown_result.success + ): + self._commit( + replace( + self._context, + shutdown_result=ShutdownResult( + False, "auxiliary runtime thread unresolved before shutdown deadline" + ), + ) + ) + self._events.put(None) + remaining = max(0.0, deadline - time.monotonic()) + self._owner.join(remaining) + owner_stopped = not self._owner.is_alive() + result = self._context.shutdown_result or ShutdownResult( + False, "shutdown result unavailable" + ) + if not owner_stopped: + result = ShutdownResult(False, "owner thread did not stop before shutdown deadline") + self._commit(replace(self._context, shutdown_result=result)) + return CommandResult(False, result, snapshot=self.snapshot()) + if not result.success: + return CommandResult(False, result, snapshot=self.snapshot()) + self._commit( + replace(self._context, shutdown=ShutdownState.CLOSED, shutdown_result=result) + ) + self._closed = True + return CommandResult(result.success and owner_stopped, result, snapshot=self.snapshot()) + + close = shutdown diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 1bb55dbc89..b4a157765e 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -25,11 +25,11 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from enum import Enum +from dataclasses import dataclass import math import threading import time -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import Any, TypeAlias from pydantic import Field @@ -39,6 +39,17 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In +from dimos.manipulation.execution_runtime import ( + ControlCoordinatorGateway, + ExecutionRuntime, + ExecutionTopology, + LifecycleState, + OperationHandle, + Outcome, + PreparedPlan, + ShutdownState, + prepare_generated_plan, +) from dimos.manipulation.planning.factory import ( KinematicsName, PlannerName, @@ -85,12 +96,8 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint from dimos.utils.logging_config import setup_logger -if TYPE_CHECKING: - from dimos.core.rpc_client import RPCClient - logger = setup_logger() # Composite type aliases for readability (using semantic IDs from planning.spec) @@ -106,14 +113,14 @@ RobotInfoPayload: TypeAlias = dict[str, RobotInfoValue] -class ManipulationState(Enum): - """State machine for manipulation module.""" +@dataclass(frozen=True) +class ManipulationExecutionSnapshot: + """Atomic execution state exposed to manipulation operators and UIs.""" - IDLE = 0 - PLANNING = 1 - EXECUTING = 2 - COMPLETED = 3 - FAULT = 4 + state: str + diagnostic: str + ready_plan_id: str | None + has_ready_plan: bool class ManipulationModuleConfig(ModuleConfig): @@ -145,6 +152,7 @@ class ManipulationModule(Module): """ config: ManipulationModuleConfig + _runtime_factory = ExecutionRuntime # Input: Joint state from coordinator (for world sync) coordinator_joint_state: In[JointState] @@ -152,18 +160,7 @@ class ManipulationModule(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - # State machine - self._state = ManipulationState.IDLE self._lock = threading.Lock() - # State transitions are local, but coordinator calls must be serialized - # with cancellation. In particular, do not let cancel race between - # two arms of a multi-arm dispatch. - self._execution_dispatch_lock = threading.Lock() - self._execution_generation = 0 - self._possibly_active_tasks: set[str] = set() - self._cancellation_in_progress = False - self._error_message = "" - self._planning_epoch = 0 # Planning components (initialized in start()) self._world_monitor: WorldMonitor | None = None @@ -173,12 +170,8 @@ def __init__(self, **kwargs: Any) -> None: # Robot registry: maps robot_name -> (world_robot_id, config, trajectory_gen) self._robots: RobotRegistry = {} - # Canonical generated plan for plan/preview/execute workflow. - # Robot-local paths and trajectories are derived from this plan on demand. - self._last_plan: GeneratedPlan | None = None - - # Coordinator integration (lazy initialized) - self._coordinator_client: RPCClient | None = None + self._execution_runtime: ExecutionRuntime | None = None + self._execution_topology: ExecutionTopology | None = None # Init joints: captured from first joint state per robot, used by go_init self._init_joints: dict[RobotName, JointState] = {} @@ -189,32 +182,59 @@ def __init__(self, **kwargs: Any) -> None: logger.info("ManipulationModule initialized") - def _ensure_execution_transaction_state(self) -> None: - """Initialize transaction fields for lightweight test harness subclasses.""" - if not hasattr(self, "_execution_dispatch_lock"): - self._execution_dispatch_lock = threading.Lock() - if not hasattr(self, "_execution_generation"): - self._execution_generation = 0 - if not hasattr(self, "_possibly_active_tasks"): - self._possibly_active_tasks = set() - if not hasattr(self, "_cancellation_in_progress"): - self._cancellation_in_progress = False - @rpc def start(self) -> None: """Start the manipulation module.""" super().start() - # Initialize planning stack - self._initialize_planning() - - # Subscribe to joint state via port - if self.coordinator_joint_state is not None: - self.coordinator_joint_state.subscribe(self._on_joint_state) - logger.info("Subscribed to coordinator_joint_state port") + try: + # Keep subscription and all resource creation in one transaction. A + # subscription can fail just as readily as planner construction. + self._initialize_planning() + if self.coordinator_joint_state is not None: + self.coordinator_joint_state.subscribe(self._on_joint_state) + logger.info("Subscribed to coordinator_joint_state port") + except Exception: + logger.exception("Failed to initialize manipulation planning") + self._unwind_startup() + raise logger.info("ManipulationModule started") + def _unwind_startup(self) -> None: + """Best-effort cleanup for a partially completed start transaction.""" + try: + self._stop_tf() + except Exception: + logger.exception("Failed to stop TF during startup unwind") + try: + if self._world_monitor is not None: + self._world_monitor.stop_all_monitors() + except Exception: + logger.exception("Failed to stop world monitor during startup unwind") + try: + runtime = self._execution_runtime + if runtime is not None: + shutdown = runtime.shutdown(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if ( + shutdown.snapshot is not None + and shutdown.snapshot.shutdown == ShutdownState.CLOSED + ): + self._execution_runtime = None + except Exception: + logger.exception("Failed to shut down runtime during startup unwind") + finally: + try: + super().stop() + except Exception: + logger.exception("Failed to stop superclass during startup unwind") + + def _stop_tf(self) -> None: + self._tf_stop_event.set() + if self._tf_thread is not None: + self._tf_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + self._tf_thread = None + def _initialize_planning(self) -> None: """Initialize world, planner, and trajectory generator.""" if not self.config.robots: @@ -253,6 +273,20 @@ def _initialize_planning(self) -> None: self._world_monitor.finalize() + self._execution_topology = ExecutionTopology.from_robot_configs(self.config.robots) + + def gateway_factory() -> ControlCoordinatorGateway: + from dimos.control.coordinator import ControlCoordinator + from dimos.core.rpc_client import RPCClient + + return ControlCoordinatorGateway(RPCClient(None, ControlCoordinator)) + + self._execution_runtime = self._runtime_factory( + gateway_factory, + topology=self._execution_topology, + action_timeout=self.config.planning_timeout, + ) + # Add floor obstacle to prevent trajectories below the table surface if self.config.floor_z is not None: fz = self.config.floor_z @@ -411,7 +445,27 @@ def _tf_publish_loop(self) -> None: @rpc def get_state(self) -> str: """Get current manipulation state name.""" - return self._state.name + if self._execution_runtime is None: + return LifecycleState.IDLE.name + return self._execution_runtime.snapshot().state.name + + @rpc + def get_execution_snapshot(self) -> ManipulationExecutionSnapshot: + """Return one atomic runtime execution snapshot for observers.""" + if self._execution_runtime is None: + return ManipulationExecutionSnapshot( + state=LifecycleState.IDLE.name, + diagnostic="", + ready_plan_id=None, + has_ready_plan=False, + ) + snapshot = self._execution_runtime.snapshot() + return ManipulationExecutionSnapshot( + state=snapshot.state.name, + diagnostic=snapshot.diagnostic or "", + ready_plan_id=snapshot.ready_plan_id, + has_ready_plan=snapshot.ready_plan is not None, + ) @rpc def get_error(self) -> str: @@ -420,51 +474,25 @@ def get_error(self) -> str: Returns: Error message or empty string """ - return self._error_message + if self._execution_runtime is None: + return "" + snapshot = self._execution_runtime.snapshot() + return snapshot.diagnostic or "" @rpc def cancel(self) -> bool: """Cancel current motion or invalidate an in-progress plan.""" - self._ensure_execution_transaction_state() - # Publish the cancellation gate before waiting for dispatch. An - # execution that arrives while cancellation is in flight must not - # become a new transaction behind the cancellation barrier. - with self._lock: - self._cancellation_in_progress = True - # Holding this lock waits for an in-flight execute RPC to resolve and - # prevents another arm from being dispatched after cancellation. - try: - with self._execution_dispatch_lock: - with self._lock: - task_names = set(self._possibly_active_tasks) - is_planning = self._state == ManipulationState.PLANNING - is_executing = self._state == ManipulationState.EXECUTING - if not (is_planning or is_executing or task_names): - return False - self._execution_generation += 1 - if is_planning: - self._planning_epoch += 1 - plan = self._last_plan - - unresolved = self._cancel_tasks(task_names) - with self._lock: - if unresolved: - self._possibly_active_tasks = unresolved - self._state = ManipulationState.FAULT - self._error_message = "Failed to cancel coordinator tasks: " + ", ".join( - sorted(unresolved) - ) - logger.error(self._error_message) - return False - self._possibly_active_tasks.clear() - self._state = ManipulationState.IDLE - finally: - with self._lock: - self._cancellation_in_progress = False - if plan is not None: - self._dismiss_preview(plan.group_ids) - logger.info("Motion cancelled") - return True + if self._execution_runtime is None: + return False + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + return self._execution_runtime.cancel_planning().accepted + if snapshot.state == LifecycleState.READY: + result = self._execution_runtime.clear_ready_plan() + if result.accepted and isinstance(snapshot.ready_plan, PreparedPlan): + self._dismiss_preview(snapshot.ready_plan.generated_plan.group_ids) + return result.accepted + return self._execution_runtime.cancel().accepted @rpc @skill @@ -474,16 +502,42 @@ def reset(self) -> SkillResult[ManipulationSkillError]: Use this after an error or fault to allow new commands. Cannot reset while a motion is executing — cancel first. """ - with self._lock: - if self._state == ManipulationState.EXECUTING: - return SkillResult.fail( - "INVALID_STATE", - "Cannot reset while executing — cancel the motion first", - ) - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - self._error_message = "" + if self._execution_runtime is None: + return SkillResult.fail("INVALID_STATE", "Execution runtime is not started") + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + result = self._execution_runtime.cancel_planning() + return ( + SkillResult.ok("Reset to IDLE — ready for new commands") + if result.accepted + else SkillResult.fail("INVALID_STATE", result.diagnostic) + ) + if snapshot.state == LifecycleState.READY: + result = self._execution_runtime.clear_ready_plan() + if not result.accepted: + return SkillResult.fail("INVALID_STATE", result.diagnostic) + if isinstance(snapshot.ready_plan, PreparedPlan): + self._dismiss_preview(snapshot.ready_plan.generated_plan.group_ids) + return SkillResult.ok("Reset to IDLE — ready for new commands") + if snapshot.state in ( + LifecycleState.DISPATCHING, + LifecycleState.RUNNING, + LifecycleState.CANCELLING, + ): + return SkillResult.fail("INVALID_STATE", "Cannot reset while execution is active") + if snapshot.state != LifecycleState.FAULT: + return SkillResult.ok("Reset to IDLE — ready for new commands") + reset_result = self._execution_runtime.reset() + if not reset_result.accepted or reset_result.value is None: + return SkillResult.fail("INVALID_STATE", reset_result.diagnostic) + # Runtime owns the reset deadline (action timeout); wait slightly past + # it so the deadline event has a chance to publish its result. + runtime_timeout = float(getattr(self._execution_runtime, "_timeout", 0.05)) + reconciled = self._execution_runtime.wait_for_reset( + reset_result.value, timeout=max(runtime_timeout, 0.05) + 0.1 + ) + if not reconciled.accepted or reconciled.value is None or not reconciled.value.success: + return SkillResult.fail("RESET_FAILED", reconciled.diagnostic or "Reset failed") return SkillResult.ok("Reset to IDLE — ready for new commands") @rpc @@ -528,41 +582,17 @@ def is_collision_free(self, joints: list[float], robot_name: RobotName | None = return self._world_monitor.is_state_valid(robot_id, joint_state) return False - def _begin_planning( - self, robot_name: RobotName | None = None - ) -> tuple[RobotName, WorldRobotID] | None: - """Check state and begin planning. Returns (robot_name, robot_id) or None. - - Args: - robot_name: Robot to plan for (required if multiple robots configured) - """ - if self._world_monitor is None: - logger.error("Planning not initialized") - return None - if (robot := self._get_robot(robot_name)) is None: - return None - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - logger.warning(f"Cannot plan: state is {self._state.name}") - return None - self._planning_epoch += 1 - self._last_plan = None - self._state = ManipulationState.PLANNING - return robot[0], robot[1] - - def _begin_group_planning(self) -> int | None: + def _begin_group_planning(self) -> str | None: """Check state and begin planning for explicit planning-group APIs.""" - if self._world_monitor is None: + if self._world_monitor is None or self._execution_runtime is None: logger.error("Planning not initialized") return None - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - logger.warning(f"Cannot plan: state is {self._state.name}") - return None - self._planning_epoch += 1 - self._last_plan = None - self._state = ManipulationState.PLANNING - return self._planning_epoch + ready = self._execution_runtime.snapshot().ready_plan + previous = ready.generated_plan if isinstance(ready, PreparedPlan) else None + token = self._execution_runtime.start_planning() + if isinstance(token, str) and previous is not None: + self._dismiss_preview(previous.group_ids) + return token if isinstance(token, str) else None def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID: """Return the unique pose-targetable group or raise if it is ambiguous.""" @@ -576,50 +606,6 @@ def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> Plan ) return group_id - def _split_plan_trajectory_by_robot( - self, plan: GeneratedPlan - ) -> dict[RobotName, JointTrajectory] | None: - """Split a stored global selected-joint trajectory into robot-local subsets.""" - trajectories: dict[RobotName, JointTrajectory] = {} - robot_indices: dict[RobotName, list[tuple[int, str]]] = {} - for index, global_name in enumerate(plan.trajectory.joint_names): - if "/" not in global_name: - logger.error("Stored trajectory joint '%s' is not globally named", global_name) - return None - robot_name, local_name = global_name.split("/", 1) - robot = self._get_robot(robot_name) - if robot is None: - return None - _, _, config, _ = robot - if local_name not in config.joint_names: - logger.error("Stored trajectory joint '%s' is not configured", global_name) - return None - robot_indices.setdefault(robot_name, []).append((index, local_name)) - - for robot_name, indexed_names in robot_indices.items(): - indices = [index for index, _ in indexed_names] - local_names = [name for _, name in indexed_names] - points = [ - TrajectoryPoint( - time_from_start=point.time_from_start, - positions=[point.positions[index] for index in indices], - velocities=[point.velocities[index] for index in indices], - ) - for point in plan.trajectory.points - ] - trajectories[robot_name] = JointTrajectory( - joint_names=local_names, - points=points, - timestamp=plan.trajectory.timestamp, - ) - logger.info( - "Stored trajectory for %s: %d selected joints, %.3fs", - robot_name, - len(local_names), - trajectories[robot_name].duration, - ) - return trajectories - @staticmethod def _assert_finite_sequence(values: Sequence[float], label: str) -> None: for value in values: @@ -738,7 +724,7 @@ def _plan_selected_path( group_ids: tuple[PlanningGroupID, ...], start: JointState, goal: JointState, - planning_epoch: int, + planning_epoch: str, ) -> GeneratedPlan | None: """Plan over explicit planning groups and store the resulting plan.""" assert self._world_monitor and self._planner @@ -769,33 +755,24 @@ def _plan_selected_path( iterations=result.iterations, message=result.message, ) - with self._lock: - if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: - logger.info("Discarding cancelled planning result") + if self._execution_runtime is None or self._execution_topology is None: + return None + try: + prepared = prepare_generated_plan(plan, self._execution_topology) + if not self._execution_runtime.complete_planning(planning_epoch, prepared).accepted: return None - self._last_plan = plan - self._state = ManipulationState.COMPLETED + except Exception as exc: + self._fail_planning_epoch(planning_epoch, f"Failed to prepare plan: {exc}") + return None return plan - def _fail(self, msg: str) -> bool: - """Set FAULT state with error message.""" - logger.warning(msg) - with self._lock: - self._state = ManipulationState.FAULT - self._error_message = msg - return False - - def _fail_planning_epoch(self, planning_epoch: int, msg: str) -> bool: + def _fail_planning_epoch(self, planning_epoch: str, msg: str) -> bool: """Fault only the still-current planning operation.""" - with self._lock: - if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: - logger.info("Discarding cancelled planning result") - return False - logger.warning(msg) - self._last_plan = None - self._state = ManipulationState.FAULT - self._error_message = msg - return False + logger.warning(msg) + return ( + self._execution_runtime is not None + and self._execution_runtime.fail_planning(planning_epoch, msg).accepted + ) def _dismiss_preview(self, group_ids: Sequence[PlanningGroupID]) -> None: """Hide the preview ghost if the world supports it.""" @@ -922,20 +899,12 @@ def solve_ik( seed: JointState | None = None, ) -> IKResult: """Compatibility wrapper for inverse_kinematics_single().""" - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - return IKResult( - status=IKStatus.NO_SOLUTION, - message=f"Cannot solve IK while state is {self._state.name}", - ) - self._state = ManipulationState.PLANNING result = self.inverse_kinematics_single( pose, robot_name=robot_name, seed=seed, check_collision=check_collision, ) - self._state = ManipulationState.COMPLETED if result.is_success() else ManipulationState.IDLE if result.is_success(): logger.info(f"IK solved, error: {result.position_error:.4f}m") return result @@ -1012,7 +981,7 @@ def generate_plan_to_joint_targets( if self._world_monitor is None or self._planner is None: return None if not joint_targets: - self._fail("At least one joint target is required") + logger.warning("At least one joint target is required") return None group_ids = tuple( @@ -1056,7 +1025,7 @@ def generate_plan_to_pose_targets( if self._world_monitor is None or self._kinematics is None: return None if not pose_targets: - self._fail("At least one pose target is required") + logger.warning("At least one pose target is required") return None stamped_targets = { planning_group_id_from_selector(group): PoseStamped( @@ -1114,7 +1083,9 @@ def preview_plan( target_fps: float = 30.0, ) -> bool: """Preview a complete generated plan in the visualizer.""" - plan = plan or self._last_plan + if plan is None and self._execution_runtime is not None: + ready = self._execution_runtime.snapshot().ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None if plan is None or not plan.path: logger.warning("No generated plan to preview") return False @@ -1140,7 +1111,10 @@ def has_planned_path(self) -> bool: Returns: True if a path is planned and ready """ - return self._last_plan is not None and bool(self._last_plan.path) + if self._execution_runtime is None: + return False + ready = self._execution_runtime.snapshot().ready_plan + return isinstance(ready, PreparedPlan) and bool(ready.generated_plan.path) @rpc def get_visualization_url(self) -> str | None: @@ -1160,12 +1134,16 @@ def clear_planned_path(self) -> bool: Returns: True if cleared """ - with self._lock: - plan = self._last_plan - self._last_plan = None - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE + plan = None + if self._execution_runtime is not None: + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + return self._execution_runtime.cancel_planning().accepted + ready = snapshot.ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None + result = self._execution_runtime.clear_ready_plan() + if not result.accepted: + return False if plan is not None: # Preserve the group selection until the public visualization # transaction has invalidated and hidden its preview. @@ -1392,335 +1370,259 @@ def set_init_joints_to_current(self, robot_name: RobotName | None = None) -> boo ) return True - def _get_coordinator_client(self) -> RPCClient | None: - """Get or create coordinator RPC client (lazy init).""" - if not any( - c.coordinator_task_name or c.gripper_hardware_id for _, c, _ in self._robots.values() - ): - return None - if self._coordinator_client is None: - from dimos.control.coordinator import ControlCoordinator - from dimos.core.rpc_client import RPCClient - - self._coordinator_client = RPCClient(None, ControlCoordinator) - return self._coordinator_client - - def _translate_trajectory_to_coordinator( - self, - trajectory: JointTrajectory, - robot_config: RobotModelConfig, - ) -> JointTrajectory: - """Translate trajectory joint names from URDF to coordinator namespace. - - Args: - trajectory: Trajectory with URDF joint names - robot_config: Robot config with joint name mapping - - Returns: - Trajectory with coordinator joint names - """ - if not robot_config.joint_name_mapping: - return trajectory # No translation needed - - # Translate joint names - coordinator_names = [ - robot_config.get_coordinator_joint_name(j) for j in trajectory.joint_names - ] - - # Create new trajectory with translated names - # Note: duration is computed automatically from points in JointTrajectory.__init__ - return JointTrajectory( - joint_names=coordinator_names, - points=trajectory.points, - timestamp=trajectory.timestamp, - ) - - def _execute_robot_trajectory(self, robot_name: RobotName, trajectory: JointTrajectory) -> bool: - """Execute one robot's local trajectory through ControlCoordinator.""" - if (robot := self._get_robot(robot_name)) is None: - return False - resolved_name, _, config, _ = robot - - if not config.coordinator_task_name: - logger.error(f"No coordinator_task_name for '{resolved_name}'") - return False - if (client := self._get_coordinator_client()) is None: - logger.error("No coordinator client") + @rpc + def execute(self) -> bool: + """Execute the current runtime-owned prepared plan.""" + handle = self._submit_execution(None) + if handle is None or self._execution_runtime is None: return False - - translated = self._translate_trajectory_to_coordinator(trajectory, config) - logger.info( - f"Executing: task='{config.coordinator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" + dispatch = self._execution_runtime.wait_for_dispatch( + handle, timeout=self._dispatch_wait_timeout() ) - - result = client.task_invoke( - config.coordinator_task_name, "execute", {"trajectory": translated} + if not dispatch.accepted: + return self._handle_dispatch_timeout(handle, dispatch.diagnostic) == "completed" + return bool( + dispatch.accepted + and dispatch.value is not None + and dispatch.value.outcome == Outcome.ACCEPTED ) - if result: - logger.info("Trajectory accepted") - return True - return self._fail("Coordinator rejected trajectory") - - @rpc - def execute(self, robot_name: RobotName | None = None) -> bool: - """Compatibility wrapper for execute_plan().""" - return self.execute_plan(robot_name=robot_name) @rpc - def execute_plan( - self, plan: GeneratedPlan | None = None, robot_name: RobotName | None = None - ) -> bool: - """Execute a generated planning-group plan through affected trajectory tasks.""" - execution = self._begin_execution(plan) - if execution is None: - return False - previous_state, target_plan, token = execution - try: - freshness_error = self._stored_plan_freshness_error(target_plan) - if freshness_error is not None: - self._restore_execution_gate(token, previous_state, freshness_error) - return False - prepared = self._prepare_execution(target_plan, robot_name) - if isinstance(prepared, str): - self._restore_execution_gate(token, previous_state, prepared) - return False - except Exception as exc: - self._restore_execution_gate( - token, previous_state, f"Failed to prepare execution: {exc}" - ) + def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: + """Execute a complete generated plan through the runtime owner.""" + handle = self._submit_execution(plan) + if handle is None or self._execution_runtime is None: return False - return self._dispatch_prepared_plan(prepared, token) - - def _begin_execution( - self, plan: GeneratedPlan | None - ) -> tuple[ManipulationState, GeneratedPlan, int] | None: - """Enter EXECUTING if the module is idle enough to dispatch.""" - self._ensure_execution_transaction_state() - with self._lock: - if self._cancellation_in_progress: - logger.warning("Cancellation is in progress; rejecting execution") - return None - target_plan = plan or self._last_plan - if target_plan is None or not target_plan.path or not target_plan.trajectory.points: - logger.warning("Stored plan is invalid or not executable") - self._error_message = "Stored plan is invalid or not executable" - return None - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - logger.warning("Manipulation state is not executable") - return None - previous_state = self._state - self._execution_generation += 1 - token = self._execution_generation - self._state = ManipulationState.EXECUTING - return previous_state, target_plan, token - - def _restore_execution_gate( - self, token: int, previous_state: ManipulationState, msg: str - ) -> None: - logger.warning(msg) - with self._lock: - if self._execution_generation == token and self._state == ManipulationState.EXECUTING: - self._state = previous_state - self._error_message = msg - - def _stored_plan_freshness_error( - self, plan: GeneratedPlan, tolerance: float = 1e-6 - ) -> str | None: - """Return why current selected joints no longer match the plan start.""" - if self._world_monitor is None: - return "Planning not initialized" - if not plan.trajectory.points: - return "Stored trajectory has no points" - first = plan.trajectory.points[0] - names = list(plan.trajectory.joint_names) - if len(names) != len(first.positions): - return "Stored trajectory start has malformed joint positions" - try: - current = self._world_monitor.current_global_joint_state() - except (AttributeError, TypeError, ValueError): - return "Current planned joints are malformed" - if not isinstance(current, JointState) or len(current.name) != len(current.position): - return "Current planned joints are malformed" - values: dict[str, float] = {} - ordered_subset: list[tuple[str, float]] = [] - planned = set(names) - for raw_name, raw_position in zip(current.name, current.position, strict=True): - name = str(raw_name) - if name in values: - return "Current planned joints are malformed" - try: - value = float(raw_position) - except (TypeError, ValueError): - return f"Current planned joint '{name}' is malformed" - if not math.isfinite(value): - return f"Current planned joint '{name}' is malformed" - values[name] = value - if name in planned: - ordered_subset.append((name, value)) - if len(ordered_subset) != len(names): - return "Current planned joints are malformed" - for (actual_name, actual), name, expected in zip( - ordered_subset, names, first.positions, strict=True - ): - if actual_name != name: - return "Current planned joints are not in stored plan order" - if not math.isfinite(actual) or not math.isfinite(float(expected)): - return f"Current planned joint '{name}' is malformed" - if abs(actual - float(expected)) > tolerance: - return f"Current planned joint '{name}' no longer matches the stored plan start" - return None + dispatch = self._execution_runtime.wait_for_dispatch( + handle, timeout=self._dispatch_wait_timeout() + ) + if not dispatch.accepted: + return self._handle_dispatch_timeout(handle, dispatch.diagnostic) == "completed" + return bool( + dispatch.accepted + and dispatch.value is not None + and dispatch.value.outcome == Outcome.ACCEPTED + ) - def _prepare_execution( - self, - plan: GeneratedPlan, - robot_name: RobotName | None = None, - ) -> tuple[tuple[RobotName, RobotModelConfig, JointTrajectory], ...] | str: - """Resolve execution inputs before dispatching any coordinator task.""" - if self._world_monitor is None: - return "Planning not initialized" - if not plan.path: - return "No generated plan" - try: - affected = list(self._world_monitor.planning_groups.select(plan.group_ids).robot_names) - except Exception as exc: - return f"Failed to resolve generated plan: {exc}" - if robot_name is not None: - if robot_name not in affected: - return f"No planned path for '{robot_name}'" - affected = [robot_name] - trajectories = self._split_plan_trajectory_by_robot(plan) - if trajectories is None: - return "Failed to split stored trajectory" - if self._get_coordinator_client() is None: - return "No coordinator client" - prepared: list[tuple[RobotName, RobotModelConfig, JointTrajectory]] = [] - for name in affected: - robot = self._get_robot(name) - if robot is None: - return f"Robot '{name}' not found" - _, _, config, _ = robot - if not config.coordinator_task_name: - return f"No coordinator_task_name for '{name}'" - traj = trajectories.get(name) - if traj is None: - return f"No planned trajectory for '{name}'" - prepared.append((name, config, traj)) - return tuple(prepared) - - def _dispatch_prepared_plan( - self, - prepared: Sequence[tuple[RobotName, RobotModelConfig, JointTrajectory]], - token: int, - ) -> bool: - """Dispatch already-validated robot trajectories; failures here fault.""" - self._ensure_execution_transaction_state() - with self._execution_dispatch_lock: - if (client := self._get_coordinator_client()) is None: - return self._fault_execution(token, "No coordinator client") - - for _robot_name, config, traj in prepared: - task_name = config.coordinator_task_name - if task_name is None: - return self._fault_execution(token, "No coordinator_task_name") - with self._lock: - if ( - self._execution_generation != token - or self._state != ManipulationState.EXECUTING - ): - return False - # Add before the RPC: a rejected or raised execute may - # have been accepted remotely. - self._possibly_active_tasks.add(task_name) - try: - translated = self._translate_trajectory_to_coordinator(traj, config) - logger.info( - "Executing: task='%s', %d pts, %.2fs", - task_name, - len(translated.points), - translated.duration, - ) - result = client.task_invoke(task_name, "execute", {"trajectory": translated}) - except Exception as exc: - return self._fault_execution(token, f"Failed to dispatch trajectory: {exc}") - if not result: - return self._fault_execution(token, "Coordinator rejected trajectory") - - logger.info("Trajectory accepted") - with self._lock: - if ( - self._execution_generation != token - or self._state != ManipulationState.EXECUTING - ): - return False - self._state = ManipulationState.COMPLETED - return True - - def _cancel_tasks(self, task_names: set[str]) -> set[str]: - """Cancel possibly active tasks, returning names whose safety is unknown.""" - if not task_names: - return set() - client = self._get_coordinator_client() - if client is None: - return set(task_names) - unresolved: set[str] = set() - for task_name in task_names: + def _submit_execution(self, plan: GeneratedPlan | None) -> OperationHandle | None: + """Submit one explicit or exactly-ready generated plan to the runtime.""" + runtime = self._execution_runtime + topology = self._execution_topology + if runtime is None or topology is None: + return None + ready = runtime.snapshot().ready_plan + if plan is None or (isinstance(ready, PreparedPlan) and ready.generated_plan is plan): + result = runtime.execute_ready() + else: + previous = ready.generated_plan if isinstance(ready, PreparedPlan) else None try: - result = client.task_invoke(task_name, "cancel", {}) - # Coordinator cancel returns False when the task is already - # non-executing; that is a confirmed-safe result. - if result is None: - unresolved.add(task_name) - except Exception: - unresolved.add(task_name) - return unresolved - - def _fault_execution(self, token: int, message: str) -> bool: - """Rollback attempted tasks and fault, without overwriting cancellation.""" - self._ensure_execution_transaction_state() - with self._lock: - if self._execution_generation != token: - return False - task_names = set(self._possibly_active_tasks) - unresolved = self._cancel_tasks(task_names) - with self._lock: - if self._execution_generation != token: - return False - if unresolved: - self._possibly_active_tasks = unresolved - message += "; unresolved coordinator tasks: " + ", ".join(sorted(unresolved)) - else: - self._possibly_active_tasks.clear() - self._state = ManipulationState.FAULT - self._error_message = message - logger.error(message) - return False - - @rpc - def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: - """Get trajectory execution status via coordinator task_invoke.""" - last_plan = self._last_plan - if robot_name is None and last_plan is not None and last_plan.path: - if self._world_monitor is None: + prepared = prepare_generated_plan(plan, topology) + except Exception as exc: + logger.warning("Failed to prepare execution: %s", exc) return None - statuses = { - name: self.get_trajectory_status(name) - for name in self._world_monitor.planning_groups.select( - last_plan.group_ids - ).robot_names - } - return {"robots": statuses} - if (robot := self._get_robot(robot_name)) is None: - return None - _, _, config, _ = robot - if not config.coordinator_task_name or (client := self._get_coordinator_client()) is None: + result = runtime.execute_explicit(prepared) + if result.accepted and previous is not None: + self._dismiss_preview(previous.group_ids) + return result.value if result.accepted else None + + def _dispatch_wait_timeout(self) -> float: + """Return the runtime dispatch deadline plus a publication margin.""" + runtime = self._execution_runtime + configured = self.config.planning_timeout + action_timeout = float(getattr(runtime, "_timeout", configured)) if runtime else configured + return max(action_timeout, 0.05) + 0.1 + + def _handle_dispatch_timeout(self, handle: OperationHandle, diagnostic: str) -> str: + """Cancel the timed-out operation and verify its terminal safety state.""" + runtime = self._execution_runtime + if runtime is None: + logger.error("Execution dispatch timed out; runtime is unavailable") + return "unresolved" + + # A timeout can race with result publication. Reconcile the original + # handle before inspecting the active operation, so a replacement can + # never be cancelled on behalf of this request. + retained_dispatch = runtime.wait_for_dispatch(handle, timeout=0.01) + retained_terminal = runtime.wait_for_terminal(handle, timeout=0.01) + terminal_result = ( + retained_terminal.value + if retained_terminal.accepted and retained_terminal.value is not None + else None + ) + if terminal_result is not None: + if terminal_result.outcome == Outcome.CANCELLED: + logger.warning("Execution dispatch timed out; operation was safely cancelled") + return "cancelled" + if terminal_result.outcome == Outcome.COMPLETED: + logger.info("Execution dispatch timed out; operation completed") + return "completed" + if terminal_result.outcome == Outcome.REJECTED: + logger.error("Execution dispatch timed out; execution was rejected") + return "rejected" + if terminal_result.outcome == Outcome.FAILED: + logger.error("Execution dispatch timed out; execution failed") + return "failed" + logger.error( + "Execution dispatch timed out; remote safety is unresolved (%s)", + terminal_result.outcome.value, + ) + return "unresolved" + if ( + retained_dispatch.accepted + and retained_dispatch.value is not None + and retained_dispatch.value.outcome != Outcome.ACCEPTED + ): + if retained_dispatch.value.outcome == Outcome.CANCELLED: + logger.warning("Execution dispatch timed out; operation was safely cancelled") + return "cancelled" + if retained_dispatch.value.outcome == Outcome.COMPLETED: + logger.info("Execution dispatch timed out; operation completed") + return "completed" + if retained_dispatch.value.outcome == Outcome.REJECTED: + logger.error("Execution dispatch timed out; execution was rejected") + return "rejected" + if retained_dispatch.value.outcome == Outcome.FAILED: + logger.error("Execution dispatch timed out; execution failed") + return "failed" + logger.error( + "Execution dispatch timed out; remote safety is unresolved (%s)", + retained_dispatch.value.outcome.value, + ) + return "unresolved" + + cancellation = runtime.cancel_if_current(handle) + if not cancellation.accepted: + concurrent_terminal = runtime.wait_for_terminal(handle, timeout=0.01) + if concurrent_terminal.accepted and concurrent_terminal.value is not None: + outcome = concurrent_terminal.value.outcome + if outcome == Outcome.CANCELLED: + logger.warning("Execution dispatch timed out; operation was safely cancelled") + return "cancelled" + if outcome == Outcome.COMPLETED: + logger.info("Execution dispatch timed out; operation completed") + return "completed" + if outcome == Outcome.REJECTED: + logger.error("Execution dispatch timed out; execution was rejected") + return "rejected" + if outcome == Outcome.FAILED: + logger.error("Execution dispatch timed out; execution failed") + return "failed" + logger.error( + "Execution dispatch timed out; remote safety is unresolved (%s)", + outcome.value, + ) + return "unresolved" + active = runtime.snapshot().operation + if active is None or active.handle != handle: + message = ( + "Execution dispatch timed out; original operation is unresolved and " + "active operation handle changed" + if active is not None + else "Execution dispatch timed out; original operation is unresolved and no active operation remains" + ) + logger.error("%s (%s)", message, diagnostic or "timeout") + return "unresolved" + message = ( + cancellation.diagnostic or "Execution dispatch timed out; cancellation was rejected" + ) + logger.error("%s", message) + return "cancel_rejected" + + terminal = runtime.wait_for_terminal(handle, timeout=self._dispatch_wait_timeout()) + if terminal.accepted and terminal.value is not None: + if terminal.value.outcome == Outcome.CANCELLED: + logger.warning("Execution dispatch timed out; operation was safely cancelled") + return "cancelled" + if terminal.value.outcome == Outcome.COMPLETED: + logger.info("Execution dispatch timed out; operation completed") + return "completed" + if terminal.value.outcome == Outcome.REJECTED: + logger.error("Execution dispatch timed out; execution was rejected") + return "rejected" + if terminal.value.outcome == Outcome.FAILED: + logger.error("Execution dispatch timed out; execution failed") + return "failed" + logger.error( + "Execution dispatch timed out; remote safety is unresolved (%s)", + terminal.value.outcome.value, + ) + return "unresolved" + else: + message = ( + "Execution dispatch timed out; cancellation was accepted but operation " + "remains unresolved" + ) + logger.error( + "%s%s", + message, + f" ({terminal.diagnostic})" if terminal.diagnostic else "", + ) + return "unresolved" + + def _terminal_after_cancel_rejection( + self, handle: OperationHandle, diagnostic: str + ) -> SkillResult[ManipulationSkillError] | None: + """Reconcile a cancellation race using only the original handle.""" + runtime = self._execution_runtime + if runtime is None: return None - try: - state = client.task_invoke(config.coordinator_task_name, "get_state", {}) - if state is not None: - return {"state": int(state), "task": config.coordinator_task_name} + terminal = runtime.wait_for_terminal(handle, timeout=0.01) + if not terminal.accepted or terminal.value is None: return None - except Exception: + result = terminal.value + if result.outcome == Outcome.COMPLETED: + return SkillResult.ok("Execution completed while cancellation was requested") + if result.outcome == Outcome.CANCELLED: + return SkillResult.fail( + "EXECUTION_CANCELLED", result.diagnostic or "Trajectory cancelled" + ) + if result.outcome == Outcome.REJECTED: + return SkillResult.fail( + "EXECUTION_REJECTED", result.diagnostic or "Trajectory rejected" + ) + if result.outcome == Outcome.FAILED: + return SkillResult.fail("EXECUTION_FAILED", result.diagnostic or "Trajectory failed") + return SkillResult.fail( + "EXECUTION_CANCEL_UNRESOLVED", + result.diagnostic or diagnostic or "Remote execution safety is unresolved", + ) + + @rpc + def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: + """Return a snapshot-only runtime status, optionally scoped to a robot.""" + if self._execution_runtime is None: return None + snapshot = self._execution_runtime.snapshot() + operation = snapshot.operation + tasks: list[dict[str, Any]] = [] + if operation is not None: + for task in operation.tasks: + task_robot = task.entry.robot_name + if task_robot is None and self._execution_topology is not None: + task_robot, _ = self._execution_topology.route_for_group( + task.entry.planning_group + ) + if robot_name is None or task_robot == robot_name: + tasks.append( + { + "task_id": task.task_id, + "task_name": task.task_name, + "robot_name": task_robot, + "activity": task.activity.value, + "cancel_required": task.cancel_required, + "reset_required": task.reset_required, + } + ) + return { + "state": snapshot.state.name, + "diagnostic": snapshot.diagnostic, + "fault": snapshot.fault, + "ready_plan_id": snapshot.ready_plan_id, + "operation_id": ( + snapshot.operation.handle.operation_id if snapshot.operation else None + ), + "tasks": tasks, + } @property def world_monitor(self) -> WorldMonitor | None: @@ -1792,11 +1694,10 @@ def _set_gripper_position(self, position: float, robot_name: RobotName | None = hw_id = self._get_gripper_hardware_id(robot_name) if hw_id is None: return False - client = self._get_coordinator_client() - if client is None: - logger.error("No coordinator client for gripper control") + if self._execution_runtime is None: return False - return bool(client.set_gripper_position(hw_id, position)) + result = self._execution_runtime.set_gripper_position(hw_id, position) + return bool(result.accepted and result.value == Outcome.ACCEPTED) @rpc def get_gripper(self, robot_name: RobotName | None = None) -> float | None: @@ -1808,11 +1709,10 @@ def get_gripper(self, robot_name: RobotName | None = None) -> float | None: hw_id = self._get_gripper_hardware_id(robot_name) if hw_id is None: return None - client = self._get_coordinator_client() - if client is None: + if self._execution_runtime is None: return None - result = client.get_gripper_position(hw_id) - return float(result) if result is not None else None + result = self._execution_runtime.get_gripper_position(hw_id) + return result.value if result.accepted else None @skill def set_gripper( @@ -1850,83 +1750,6 @@ def close_gripper(self, robot_name: str | None = None) -> SkillResult[Manipulati return SkillResult.ok("Gripper closed") return SkillResult.fail("GRIPPER_FAILED", "Failed to close gripper") - def _wait_for_trajectory_completion( - self, robot_name: RobotName | None = None, timeout: float = 60.0, poll_interval: float = 0.2 - ) -> bool: - """Wait for trajectory execution to complete. - - Polls the coordinator task state via task_invoke. Falls back to waiting - for the trajectory duration if the coordinator is unavailable. - - Args: - robot_name: Robot to monitor - timeout: Maximum wait time in seconds - poll_interval: Time between status checks - - Returns: - True if trajectory completed successfully - """ - last_plan = self._last_plan - if robot_name is None and last_plan is not None and last_plan.path: - try: - assert self._world_monitor is not None - affected = self._world_monitor.planning_groups.select( - last_plan.group_ids - ).robot_names - except Exception as exc: - logger.warning("Failed to resolve generated plan while waiting: %s", exc) - return False - return all( - self._wait_for_trajectory_completion(name, timeout, poll_interval) - for name in affected - ) - - robot = self._get_robot(robot_name) - if robot is None: - return True - _, _, config, _ = robot - client = self._get_coordinator_client() - trajectory = last_plan.trajectory if last_plan is not None else None - - if client is None or not config.coordinator_task_name: - # No coordinator — wait for trajectory duration as fallback - if trajectory is not None: - logger.info( - f"No coordinator status — waiting {trajectory.duration:.1f}s for trajectory" - ) - time.sleep(trajectory.duration + 0.5) - return True - - # Poll task state via task_invoke - start = time.time() - while (time.time() - start) < timeout: - try: - state = client.task_invoke(config.coordinator_task_name, "get_state", {}) - # TrajectoryState is an IntEnum: IDLE=0, EXECUTING=1, COMPLETED=2, ABORTED=3, FAULT=4 - if state is not None: - state_val = int(state) - if state_val in (0, 2): # IDLE or COMPLETED - return True - if state_val in (3, 4): # ABORTED or FAULT - logger.warning(f"Trajectory failed: state={state}") - return False - # state_val == 1 means EXECUTING, keep polling - else: - # task_invoke returned None — task not found, assume done - return True - except Exception: - # Fallback: wait for trajectory duration - if trajectory is not None: - remaining = trajectory.duration - (time.time() - start) - if remaining > 0: - logger.info(f"Status poll failed — waiting {remaining:.1f}s for trajectory") - time.sleep(remaining + 0.5) - return True - time.sleep(poll_interval) - - logger.warning(f"Trajectory execution timed out after {timeout}s") - return False - def _lift_if_low( self, robot_name: RobotName | None = None, min_z: float = 0.05 ) -> SkillResult[ManipulationSkillError]: @@ -1958,11 +1781,107 @@ def _preview_execute_wait( self.preview_path(preview_duration, robot_name) logger.info("Executing trajectory...") - if not self.execute(robot_name): + if self._execution_runtime is None: + return SkillResult.fail("EXECUTION_FAILED", "Execution runtime is not started") + ready = self._execution_runtime.snapshot().ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None + handle = self._submit_execution(plan) + if handle is None: return SkillResult.fail("EXECUTION_FAILED", "Trajectory execution failed") - if not self._wait_for_trajectory_completion(robot_name): - return SkillResult.fail("EXECUTION_TIMEOUT", "Trajectory execution timed out") + dispatch = self._execution_runtime.wait_for_dispatch( + handle, timeout=self._dispatch_wait_timeout() + ) + if not dispatch.accepted: + safety = self._handle_dispatch_timeout(handle, dispatch.diagnostic) + if safety == "cancel_rejected": + return SkillResult.fail( + "EXECUTION_CANCEL_REJECTED", + dispatch.diagnostic + or "Execution dispatch timed out; cancellation was rejected", + ) + if safety == "cancelled": + return SkillResult.fail( + "EXECUTION_CANCELLED", + "Execution dispatch timed out; operation was safely cancelled", + ) + if safety == "completed": + return SkillResult.ok("Execution dispatch timed out; operation completed") + if safety == "failed": + return SkillResult.fail( + "EXECUTION_FAILED", + "Execution dispatch timed out; operation resolved with failure", + ) + if safety == "rejected": + return SkillResult.fail( + "EXECUTION_REJECTED", "Execution dispatch timed out; execution was rejected" + ) + return SkillResult.fail( + "EXECUTION_CANCEL_UNRESOLVED", + "Execution dispatch timed out; cancellation was accepted but operation is unresolved", + ) + if dispatch.value is None: + return SkillResult.fail( + "EXECUTION_DISPATCH_FAILED", "Execution dispatch returned no result" + ) + if dispatch.value.outcome == Outcome.REJECTED: + return SkillResult.fail( + "EXECUTION_REJECTED", dispatch.value.diagnostic or "Trajectory rejected" + ) + if dispatch.value.outcome != Outcome.ACCEPTED: + return SkillResult.fail( + "EXECUTION_DISPATCH_FAILED", + dispatch.value.diagnostic + or f"Unexpected dispatch outcome: {dispatch.value.outcome.value}", + ) + + duration = plan.trajectory.duration if plan is not None else 0.0 + execution_timeout = max(float(duration) + 1.0, 0.5) + terminal = self._execution_runtime.wait_for_terminal(handle, timeout=execution_timeout) + if not terminal.accepted: + cancellation = self._execution_runtime.cancel_if_current(handle) + if not cancellation.accepted: + reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) + if reconciled is not None: + return reconciled + return SkillResult.fail( + "EXECUTION_CANCEL_REJECTED", + cancellation.diagnostic or "Execution timed out and cancellation was rejected", + ) + return SkillResult.fail( + "EXECUTION_CANCEL_UNRESOLVED", + "Execution timed out; cancellation was accepted but motion is unresolved", + ) + if terminal.value is None: + cancellation = self._execution_runtime.cancel_if_current(handle) + if not cancellation.accepted: + reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) + if reconciled is not None: + return reconciled + return SkillResult.fail("EXECUTION_FAILED", "Execution returned no result") + outcome = terminal.value.outcome + if outcome == Outcome.REJECTED: + return SkillResult.fail( + "EXECUTION_REJECTED", terminal.value.diagnostic or "Trajectory rejected" + ) + if outcome == Outcome.FAILED: + return SkillResult.fail( + "EXECUTION_FAILED", terminal.value.diagnostic or "Trajectory failed" + ) + if outcome == Outcome.CANCELLED: + return SkillResult.fail( + "EXECUTION_CANCELLED", terminal.value.diagnostic or "Trajectory cancelled" + ) + if outcome != Outcome.COMPLETED: + cancellation = self._execution_runtime.cancel_if_current(handle) + if not cancellation.accepted: + reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) + if reconciled is not None: + return reconciled + return SkillResult.fail( + "EXECUTION_CANCEL_UNRESOLVED", + f"Remote execution safety is unresolved: {outcome.value}", + ) return SkillResult.ok() @@ -2206,14 +2125,31 @@ def stop(self) -> None: """Stop the manipulation module.""" logger.info("Stopping ManipulationModule") - # Stop TF thread - if self._tf_thread is not None: - self._tf_stop_event.set() - self._tf_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self._tf_thread = None - - # Stop world monitor (includes visualization thread) - if self._world_monitor is not None: - self._world_monitor.stop_all_monitors() - - super().stop() + try: + try: + runtime = self._execution_runtime + if runtime is not None: + shutdown = runtime.shutdown(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if ( + shutdown.snapshot is not None + and shutdown.snapshot.shutdown == ShutdownState.CLOSED + ): + self._execution_runtime = None + except Exception: + logger.exception("Failed to shut down execution runtime") + finally: + try: + self._stop_tf() + except Exception: + logger.exception("Failed to stop TF thread") + finally: + try: + if self._world_monitor is not None: + self._world_monitor.stop_all_monitors() + except Exception: + logger.exception("Failed to stop world monitor") + finally: + try: + super().stop() + except Exception: + logger.exception("Failed to stop superclass") diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index 6bc7582239..5c9718ecb1 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -26,7 +26,7 @@ joints() Get current joint positions ee() Get end-effector pose groups() List explicit planning groups - state() Get module state (IDLE, PLANNING, EXECUTING, ...) + state() Get module state (IDLE, PLANNING, RUNNING, ...) ik_pose(x,y,z, seed_joints=None) Solve IK only, without path planning ik_group_pose(group_id,x,y,z) Solve IK for an explicit planning group plan(joints) Plan to joint configuration, e.g. plan([0.1]*7) @@ -221,9 +221,9 @@ def preview( return _client.preview_plan(None, duration, robot_name) -def execute(robot_name: str | None = None) -> bool: +def execute() -> bool: """Execute planned trajectory via coordinator.""" - return _client.execute(robot_name) + return _client.execute() def home(robot_name: str | None = None) -> bool: @@ -233,7 +233,7 @@ def home(robot_name: str | None = None) -> bool: home_joints = _client.get_robot_info(robot_name).get("home_joints", [0.0] * 7) success = _client.plan_to_joints(JointState(position=home_joints), robot_name) if success: - return _client.execute(robot_name) + return _client.execute() return False diff --git a/dimos/manipulation/skill_errors.py b/dimos/manipulation/skill_errors.py index 9a17085ec5..0c4e4a4b63 100644 --- a/dimos/manipulation/skill_errors.py +++ b/dimos/manipulation/skill_errors.py @@ -37,6 +37,12 @@ "GRASP_ATTEMPTS_EXHAUSTED", "GRIPPER_FAILED", "WORLD_MONITOR_UNAVAILABLE", + "EXECUTION_CANCELLED", + "EXECUTION_REJECTED", + "EXECUTION_CANCEL_REJECTED", + "EXECUTION_CANCEL_UNRESOLVED", + "EXECUTION_DISPATCH_FAILED", + "RESET_FAILED", ] # Union of codes a manipulation skill may emit (common + manipulation-specific). diff --git a/dimos/manipulation/test_execution_runtime.py b/dimos/manipulation/test_execution_runtime.py new file mode 100644 index 0000000000..1811063b17 --- /dev/null +++ b/dimos/manipulation/test_execution_runtime.py @@ -0,0 +1,1811 @@ +"""Focused owner/reconciliation tests.""" + +from dataclasses import dataclass, replace +from pathlib import Path +import threading +import time +from typing import Any + +import pytest + +from dimos.manipulation.execution_runtime import ( + ActionMethod, + ControlCoordinatorGateway, + ExecutionPlan, + ExecutionRuntime, + ExecutionTopology, + LifecycleState, + Outcome, + PreparedPlan, + ResetHandle, + ResetResult, + ShutdownState, + TaskActivity, + TaskEntry, + _Call, + _Event, + _GatewayGateState, + prepare_generated_plan, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +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.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +@dataclass +class FakeRPC: + values: dict[tuple[str, str], Any] + calls: list[tuple[str, str, dict[str, Any]]] + entered: threading.Event | None = None + release: threading.Event | None = None + + def task_invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + self.calls.append((task, method, kwargs)) + if method == "execute" and self.entered is not None: + self.entered.set() + assert self.release is not None and self.release.wait(2) + return self.values.get((task, method)) + + def stop_rpc_client(self) -> None: + self.calls.append(("", "stop", {})) + + def set_gripper_position(self, hardware_id: str, position: float) -> Any: + self.calls.append((hardware_id, "set_gripper_position", {"position": position})) + return self.values.get((hardware_id, "set_gripper_position"), True) + + def get_gripper_position(self, hardware_id: str) -> Any: + self.calls.append((hardware_id, "get_gripper_position", {})) + return self.values.get((hardware_id, "get_gripper_position"), 0.5) + + +class BlockingGateway: + """Deterministic per-task barriers for effect/correlation assertions.""" + + def __init__( + self, + executes: dict[str, Outcome], + cancels: dict[str, Outcome] | None = None, + *, + block_execute: bool = False, + ) -> None: + self.executes = executes + self.cancels = cancels or {task: Outcome.CANCELLED for task in executes} + self.block_execute = block_execute + self.calls: list[tuple[str, ActionMethod]] = [] + self.entered = {task: threading.Event() for task in executes} + self.execute_entered = {task: threading.Event() for task in executes} + self.cancel_entered = {task: threading.Event() for task in executes} + self.release = {task: threading.Event() for task in executes} + self._lock = threading.Lock() + self.stop_calls = 0 + self.calls_at_stop = 0 + + def _call(self, task: str, method: ActionMethod, outcome: Outcome) -> Outcome: + with self._lock: + self.calls.append((task, method)) + if method == ActionMethod.CANCEL: + self.entered[task].set() + self.cancel_entered[task].set() + assert self.release[task].wait(2) + return outcome + + def execute(self, task_name: str, request: Any) -> Outcome: + self.execute_entered[task_name].set() + if self.block_execute: + assert self.release[task_name].wait(2) + return self._call(task_name, ActionMethod.EXECUTE, self.executes[task_name]) + + def cancel(self, task_name: str) -> Outcome: + return self._call(task_name, ActionMethod.CANCEL, self.cancels[task_name]) + + def status(self, task_name: str) -> Outcome: + return Outcome.UNKNOWN + + def reset(self, task_name: str) -> Outcome: + return self._call(task_name, ActionMethod.RESET, Outcome.INACTIVE) + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return self._call(hardware_id, ActionMethod.GRIPPER_SET, Outcome.ACCEPTED) + + def get_gripper_position(self, hardware_id: str) -> float: + with self._lock: + self.calls.append((hardware_id, ActionMethod.GRIPPER_GET)) + return 0.25 + + def stop(self) -> None: + with self._lock: + self.stop_calls += 1 + self.calls_at_stop = len(self.calls) + for event in self.release.values(): + event.set() + + def count(self, task: str, method: ActionMethod) -> int: + with self._lock: + return self.calls.count((task, method)) + + +class BlockingAuxGateway(BlockingGateway): + def __init__(self) -> None: + super().__init__({"gripper": Outcome.ACCEPTED}) + self.aux_entered = threading.Event() + self.aux_release = threading.Event() + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + with self._lock: + self.calls.append((hardware_id, ActionMethod.GRIPPER_SET)) + self.aux_entered.set() + assert self.aux_release.wait(2) + return Outcome.ACCEPTED + + +def plan(*names: str) -> ExecutionPlan: + return ExecutionPlan( + "generated", tuple(TaskEntry(name, name, {"trajectory": name}) for name in names) + ) + + +def wait_until(fn: Any, timeout: float = 1.0) -> None: + end = time.monotonic() + timeout + while time.monotonic() < end and not fn(): + time.sleep(0.01) + assert fn() + + +def task_is(runtime: ExecutionRuntime, index: int, activity: TaskActivity) -> bool: + operation = runtime.snapshot().operation + return operation is not None and operation.tasks[index].activity == activity + + +def test_gateway_normalizes_task_invoke_and_real_states() -> None: + rpc = FakeRPC( + { + ("a", "execute"): True, + ("a", "cancel"): False, + ("a", "get_state"): TrajectoryState.EXECUTING, + }, + [], + ) + gateway = ControlCoordinatorGateway(rpc) + assert gateway.execute("a", {}) == Outcome.ACCEPTED + assert gateway.cancel("a") == Outcome.INACTIVE + assert gateway.status("a") == Outcome.RUNNING + + +def test_aggregate_dispatch_pre_registers_and_waits_for_each_acceptance() -> None: + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + result = runtime.execute_explicit(plan("a", "b")) + assert result.accepted + + def dispatch_processed() -> bool: + snapshot = runtime.snapshot() + return ( + len([call for call in rpc.calls if call[1] == "execute"]) == 2 + and snapshot.state == LifecycleState.RUNNING + and snapshot.operation is not None + and all(task.activity == TaskActivity.ACTIVE for task in snapshot.operation.tasks) + ) + + wait_until(dispatch_processed) + assert runtime.snapshot().state == LifecycleState.RUNNING + operation = runtime.snapshot().operation + assert operation is not None + assert len(operation.tasks) == 2 + finally: + runtime.shutdown() + + +def test_concurrent_execute_admission_is_serial_and_snapshot_coherent() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + barrier = threading.Barrier(3) + results: list[Any] = [] + + def submit() -> None: + barrier.wait() + results.append(runtime.execute_explicit(plan("a"))) + + threads = [threading.Thread(target=submit) for _ in range(2)] + try: + for thread in threads: + thread.start() + barrier.wait() + for thread in threads: + thread.join(1) + assert len(results) == 2 + assert sum(result.accepted for result in results) == 1 + rejected = next(result for result in results if not result.accepted) + assert rejected.value is not None + retained = dict(runtime._context.terminal_results)[rejected.value] + assert retained.outcome == Outcome.REJECTED + assert runtime.snapshot().operation is not None + assert runtime.snapshot().state in (LifecycleState.DISPATCHING, LifecycleState.RUNNING) + finally: + runtime.cancel() + runtime.shutdown() + + +def test_ready_replacement_failure_and_consumed_identity() -> None: + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) + try: + first_token = runtime.start_planning() + assert first_token is not None + assert runtime.complete_planning(first_token, plan("a")).accepted + replacement_token = runtime.start_planning() + assert replacement_token is not None + assert runtime.snapshot().ready_plan is None + assert not runtime.complete_planning( + replacement_token, ExecutionPlan("generated", ()) + ).accepted + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().ready_plan is None + + token = runtime.start_planning() + assert token is not None + assert runtime.complete_planning(token, plan("a")).accepted + executed = runtime.execute_ready() + assert executed.accepted and executed.value is not None + assert runtime.snapshot().ready_plan is None + consumed = runtime.execute_ready() + assert not consumed.accepted + assert consumed.value is None + finally: + runtime.cancel() + runtime.shutdown() + + +def test_active_attempt_admission_rejects_plan_and_execute_but_accepts_cancel() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + planning = runtime.start_planning() + assert planning is not None + assert not runtime.execute_explicit(plan("a")).accepted + assert runtime.cancel_planning().accepted + + started = runtime.execute_explicit(plan("a")) + assert started.accepted + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + assert not runtime.start_planning() + rejected = runtime.execute_explicit(plan("a")) + assert not rejected.accepted and rejected.value is not None + assert dict(runtime._context.terminal_results)[rejected.value].outcome == Outcome.REJECTED + assert runtime.cancel().accepted + assert not runtime.reset().accepted + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_rejection_compensates_prior_acceptance_and_unknown_is_fault() -> None: + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): False, ("a", "cancel"): True}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + runtime.execute_explicit(plan("a", "b")) + wait_until(lambda: any(call[1] == "cancel" for call in rpc.calls)) + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + assert dict(runtime._context.terminal_results).popitem()[1].outcome == Outcome.REJECTED + finally: + runtime.shutdown() + + +def test_cancel_during_execute_handles_late_acceptance() -> None: + entered, release = threading.Event(), threading.Event() + rpc = FakeRPC({("a", "execute"): True, ("a", "cancel"): True}, [], entered, release) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + assert entered.wait(1) + assert runtime.cancel().accepted + release.set() + wait_until(lambda: any(call[1] == "cancel" for call in rpc.calls)) + finally: + release.set() + runtime.shutdown() + + +def test_planning_tokens_and_explicit_execution_are_atomic() -> None: + rpc = FakeRPC({}, []) + runtime = ExecutionRuntime( + lambda: ControlCoordinatorGateway(rpc), + topology=ExecutionTopology((("a", "robot", "a"),)), + poll_interval=10, + ) + try: + token = runtime.start_planning() + assert token is not None + assert not runtime.complete_planning("stale", plan("a")).accepted + planned = runtime.complete_planning(token, plan("a")) + assert planned.accepted and planned.value is not None + assert runtime.snapshot().state == LifecycleState.READY + handle = runtime.execute_ready().value + assert handle is not None + assert handle.plan_id == planned.value + assert runtime.snapshot().ready_plan is None + finally: + runtime.shutdown() + + +def test_wait_for_reset_snapshot_is_constructed_outside_context_lock() -> None: + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) + reset_handle = ResetHandle("reset-attempt") + try: + reset_result = ResetResult(reset_handle, True) + runtime._commit(replace(runtime._context, reset_results=((reset_handle, reset_result),))) + result = runtime.wait_for_reset(reset_handle, timeout=1) + assert result.accepted and result.value == reset_result + finally: + runtime.shutdown() + + +def test_deadline_preserves_unresolved_action_and_stop_rpc_is_deferred() -> None: + entered, release = threading.Event(), threading.Event() + rpc = FakeRPC({("a", "execute"): True}, [], entered, release) + runtime = ExecutionRuntime( + lambda: ControlCoordinatorGateway(rpc), action_timeout=0.01, poll_interval=10 + ) + runtime.execute_explicit(plan("a")) + assert entered.wait(1) + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + assert runtime.snapshot().operation is not None + release.set() + runtime.shutdown() + assert any(call[1] == "stop" for call in rpc.calls) + + +def test_running_inactive_reconciles_peers_before_failed_terminal() -> None: + rpc = FakeRPC( + { + ("a", "execute"): True, + ("b", "execute"): True, + ("a", "get_state"): TrajectoryState.IDLE, + ("b", "get_state"): TrajectoryState.EXECUTING, + ("b", "cancel"): False, + }, + [], + ) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + runtime.execute_explicit(plan("a", "b")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + result = dict(runtime._context.terminal_results).popitem()[1] + assert result.outcome == Outcome.FAILED + assert any(task == "b" and method == "cancel" for task, method, _ in rpc.calls) + finally: + runtime.shutdown() + + +def test_status_unknown_faults_and_cleans_active_peers_without_completion() -> None: + rpc = FakeRPC( + { + ("a", "execute"): True, + ("b", "execute"): True, + ("a", "get_state"): None, + ("b", "cancel"): True, + }, + [], + ) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + runtime.execute_explicit(plan("a", "b")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + assert not runtime._context.terminal_results + wait_until(lambda: any(task == "b" and method == "cancel" for task, method, _ in rpc.calls)) + finally: + runtime.shutdown() + + +def test_dispatch_result_is_stored_once_for_acceptance_and_rejection() -> None: + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.ACCEPTED + finally: + runtime.shutdown() + + +def test_blocked_first_execute_preregisters_batch_and_delays_second() -> None: + entered, release = threading.Event(), threading.Event() + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, [], entered, release) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert entered.wait(1) + operation = runtime.snapshot().operation + assert operation is not None and len(operation.tasks) == 2 + assert all(task.action is not None for task in operation.tasks[:1]) + assert not any(task == "b" and method == "execute" for task, method, _ in rpc.calls) + assert runtime.snapshot().state == LifecycleState.DISPATCHING + release.set() + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 + finally: + release.set() + runtime.shutdown() + + +def test_three_task_rejection_skips_unstarted_task_and_rejects_terminally() -> None: + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): False, ("a", "cancel"): True}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b", "c")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + assert not any(task == "c" and method == "execute" for task, method, _ in rpc.calls) + assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 + assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.REJECTED + finally: + runtime.shutdown() + + +def test_cancel_unresolved_dispatch_skips_remaining_tasks() -> None: + entered, release = threading.Event(), threading.Event() + rpc = FakeRPC({("a", "execute"): True, ("a", "cancel"): True}, [], entered, release) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b", "c")).value + assert handle is not None + assert entered.wait(1) + runtime.cancel() + release.set() + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + assert not any(method == "execute" and task in ("b", "c") for task, method, _ in rpc.calls) + assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 + assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.CANCELLED + finally: + release.set() + runtime.shutdown() + + +def test_deadline_late_accept_and_reject_keep_fault_and_one_dispatch_result() -> None: + for late, expected_calls in ((True, 1), (False, 0)): + entered, release = threading.Event(), threading.Event() + rpc = FakeRPC( + {("a", "execute"): True if late else False, ("a", "cancel"): True}, [], entered, release + ) + runtime = ExecutionRuntime( + lambda rpc=rpc: ControlCoordinatorGateway(rpc), action_timeout=0.01, poll_interval=10 + ) + try: + handle = runtime.execute_explicit(plan("a")).value + assert entered.wait(1) + wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.FAULT) + release.set() + if late: + wait_until(lambda rpc=rpc: any(method == "cancel" for _, method, _ in rpc.calls)) + time.sleep(0.03) + assert runtime.snapshot().state == LifecycleState.FAULT + assert ( + len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 + ) + assert len([call for call in rpc.calls if call[1] == "cancel"]) == expected_calls + finally: + release.set() + runtime.shutdown() + + +def test_remote_fault_marks_task_and_cancels_active_peer_without_success() -> None: + rpc = FakeRPC( + { + ("a", "execute"): True, + ("b", "execute"): True, + ("a", "get_state"): TrajectoryState.FAULT, + ("b", "cancel"): True, + }, + [], + ) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + runtime.execute_explicit(plan("a", "b")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None + assert operation.tasks[0].activity.name == "REMOTE_FAULT" + wait_until(lambda: any(task == "b" and method == "cancel" for task, method, _ in rpc.calls)) + assert not runtime._context.terminal_results + finally: + runtime.shutdown() + + rpc = FakeRPC({("a", "execute"): False}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a")).value + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.REJECTED + finally: + runtime.shutdown() + + +def test_cleanup_allocates_distinct_cancel_actions_for_all_active_tasks() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b", "c")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + assert runtime.cancel().accepted + for task in ("a", "b", "c"): + assert gateway.entered[task].wait(1) + operation = runtime.snapshot().operation + assert operation is not None + actions = {task.task_name: task.action for task in operation.tasks} + assert all(actions[name] is not None for name in ("a", "b", "c")) + assert len({actions[name].action_id for name in ("a", "b", "c")}) == 3 # type: ignore[union-attr] + for task in ("a", "b", "c"): + gateway.release[task].set() + wait_until(lambda task=task: gateway.count(task, ActionMethod.CANCEL) == 1) + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.CANCELLED + assert all(gateway.count(task, ActionMethod.CANCEL) == 1 for task in ("a", "b", "c")) + finally: + runtime.shutdown() + + +def test_execute_unknown_faults_and_cleans_only_dispatched_execution() -> None: + gateway = BlockingGateway({"a": Outcome.UNKNOWN, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b", "c")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None and operation.uncertain + assert gateway.entered["a"].wait(1) + assert not any( + task in ("b", "c") and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + for event in gateway.release.values(): + event.set() + wait_until( + lambda: all(gateway.count(task, ActionMethod.CANCEL) == 1 for task in ("a", "b", "c")) + ) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.UNKNOWN + finally: + runtime.shutdown() + + +def test_cancelled_unresolved_dispatch_retains_cancelled_result_after_late_rejection() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}, {"a": Outcome.CANCELLED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None and gateway.execute_entered["a"].wait(1) + assert runtime.cancel().accepted + gateway.release["a"].set() + wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.CANCELLED + assert gateway.count("a", ActionMethod.CANCEL) == 0 + assert not any( + task == "b" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + finally: + runtime.shutdown() + + +def test_deadline_duplicate_and_late_result_keep_one_original_dispatch_record() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.01, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a")).value + assert handle is not None + assert gateway.execute_entered["a"].wait(1) + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None and operation.tasks[0].action is not None + task = operation.tasks[0] + action = task.action + assert action is not None + runtime._submit( + lambda: runtime._reduce( + _Event( + "action_deadline", + handle.operation_id, + handle.attempt_id, + task.task_id, + action.action_id, + action.method, + ) + ) + ) + gateway.release["a"].set() + assert gateway.cancel_entered["a"].wait(1) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.UNKNOWN + gateway.release["a"].set() + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_partial_rejection_after_three_accepts_retains_rejected_dispatch_once() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED, "d": Outcome.REJECTED} + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b", "c", "d")).value + assert handle is not None + wait_until(lambda: gateway.count("d", ActionMethod.EXECUTE) == 1) + + def rejected_dispatch_processed() -> bool: + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + return len(records) == 1 and records[0][1].outcome == Outcome.REJECTED + + wait_until(rejected_dispatch_processed) + records = [item for item in runtime._context.dispatch_results if item[0] == handle] + assert len(records) == 1 and records[0][1].outcome == Outcome.REJECTED + for event in gateway.release.values(): + event.set() + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_late_response_updates_knowledge_without_accepting_stale_action() -> None: + for outcome, expected_activity in ( + (Outcome.ACCEPTED, TaskActivity.CANCELLED), + (Outcome.REJECTED, TaskActivity.INACTIVE), + ): + gateway = BlockingGateway({"a": outcome}, block_execute=True) + runtime = ExecutionRuntime( + lambda gateway=gateway: gateway, action_timeout=0.01, poll_interval=10 + ) + try: + handle = runtime.execute_explicit(plan("a")).value + assert handle is not None and gateway.execute_entered["a"].wait(1) + wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.FAULT) + gateway.release["a"].set() + if outcome == Outcome.ACCEPTED: + wait_until(lambda gateway=gateway: gateway.count("a", ActionMethod.CANCEL) == 1) + gateway.release["a"].set() + wait_until( + lambda runtime=runtime, expected_activity=expected_activity: task_is( + runtime, 0, expected_activity + ) + ) + operation = runtime.snapshot().operation + assert operation is not None + assert operation.tasks[0].action is None + task_id = operation.tasks[0].task_id + operation_id, attempt_id = handle.operation_id, handle.attempt_id + stale = runtime._submit( + lambda runtime=runtime, + operation_id=operation_id, + attempt_id=attempt_id, + task_id=task_id: runtime._reduce( + _Event( + "action_finished", + operation_id, + attempt_id, + task_id, + "wrong", + ActionMethod.CANCEL, + Outcome.CANCELLED, + ) + ) + ) + assert stale is False + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_remote_fault_requires_reset_and_cancels_peer_without_success() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, + {"a": Outcome.CANCELLED, "b": Outcome.CANCELLED}, + ) + gateway.status = lambda task: Outcome.FAILED if task == "a" else Outcome.RUNNING # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None and operation.tasks[0].activity == TaskActivity.REMOTE_FAULT + assert operation.tasks[0].reset_required + assert gateway.cancel_entered["b"].wait(1) + gateway.release["b"].set() + wait_until(lambda: task_is(runtime, 1, TaskActivity.CANCELLED)) + operation = runtime.snapshot().operation + assert operation is not None and all(task.action is None for task in operation.tasks) + assert runtime.snapshot().state == LifecycleState.FAULT + assert not runtime._context.terminal_results + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_waiters_use_only_the_requested_result_map() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.REJECTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + assert gateway.cancel_entered["a"].wait(1) + dispatch = runtime.wait_for_dispatch(handle, timeout=0.1) + assert ( + dispatch.accepted + and dispatch.value is not None + and dispatch.value.outcome == Outcome.REJECTED + ) + terminal = runtime.wait_for_terminal(handle, timeout=0.01) + assert not terminal.accepted and terminal.diagnostic == "timeout" + gateway.release["a"].set() + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert ( + terminal.accepted + and terminal.value is not None + and terminal.value.outcome == Outcome.REJECTED + ) + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_execute_failure_latches_remote_fault_and_cleans_peer() -> None: + gateway = BlockingGateway({"a": Outcome.FAILED, "b": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None + assert operation.tasks[0].activity == TaskActivity.REMOTE_FAULT + assert operation.tasks[0].reset_required + assert gateway.cancel_entered["b"].wait(1) + assert not any( + task == "b" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + gateway.release["b"].set() + assert not runtime._context.terminal_results + finally: + gateway.release["b"].set() + runtime.shutdown() + + +def test_cancel_failure_latches_remote_fault_and_cleans_peer() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, + {"a": Outcome.FAILED, "b": Outcome.CANCELLED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + assert runtime.cancel().accepted + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + gateway.release["a"].set() + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + operation = runtime.snapshot().operation + assert operation is not None + assert operation.tasks[0].activity == TaskActivity.REMOTE_FAULT + assert operation.tasks[0].reset_required + gateway.release["b"].set() + wait_until(lambda: task_is(runtime, 1, TaskActivity.CANCELLED)) + operation = runtime.snapshot().operation + assert operation is not None and all(task.action is None for task in operation.tasks) + assert not runtime._context.terminal_results + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_result_retention_is_bounded_and_handles_remain_correlated() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + handles = [] + try: + for _ in range(20): + result = runtime.execute_explicit(plan("a")) + assert result.accepted and result.value is not None + handle = result.value + handles.append(handle) + dispatch = runtime.wait_for_dispatch(handle, timeout=1) + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.handle == handle + assert terminal.value.handle == handle + assert dispatch.value.outcome == Outcome.REJECTED + assert terminal.value.outcome == Outcome.REJECTED + assert len(runtime._context.dispatch_results) == 16 + assert len(runtime._context.terminal_results) == 16 + assert all(handle in dict(runtime._context.dispatch_results) for handle in handles[-16:]) + assert all(handle in dict(runtime._context.terminal_results) for handle in handles[-16:]) + assert handles[0] not in dict(runtime._context.dispatch_results) + assert handles[0] not in dict(runtime._context.terminal_results) + finally: + runtime.shutdown() + + +def test_unknown_fault_reset_fails_and_leaves_fault_latched() -> None: + gateway = BlockingGateway({"a": Outcome.UNKNOWN}) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.1, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + assert gateway.cancel_entered["a"].wait(1) + gateway.release["a"].set() + handle_result = runtime.reset() + assert handle_result.accepted and handle_result.value is not None + reset = runtime.wait_for_reset(handle_result.value, timeout=1) + assert reset.accepted and reset.value is not None and not reset.value.success + assert runtime.snapshot().state == LifecycleState.FAULT + time.sleep(0.05) + assert runtime.snapshot().state == LifecycleState.FAULT + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_remote_fault_reset_proves_inactive_and_resets_remote_task() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) + status_calls: dict[str, int] = {"a": 0, "b": 0} + + def status(task: str) -> Outcome: + status_calls[task] += 1 + return Outcome.FAILED if task == "a" and status_calls[task] == 1 else Outcome.INACTIVE + + gateway.status = status # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) + try: + runtime.execute_explicit(plan("a", "b")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + gateway.release["b"].set() + wait_until(lambda: gateway.count("b", ActionMethod.CANCEL) == 1) + reset_handle = runtime.reset().value + assert reset_handle is not None + reset = runtime.wait_for_reset(reset_handle, timeout=1) + assert reset.accepted and reset.value is not None and reset.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + assert gateway.count("a", ActionMethod.RESET) == 1 + assert runtime._context.reset_handle is None + assert not runtime._context.reset_proven_inactive + assert not runtime._context.reset_completed_tasks + assert dict(runtime._context.reset_results)[reset_handle] == reset.value + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_remote_reset_failure_keeps_fault_latched() -> None: + for reset_outcome in (Outcome.RUNNING, Outcome.UNKNOWN): + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + gateway.status = lambda task: Outcome.INACTIVE # type: ignore[method-assign] + gateway.reset = lambda task, outcome=reset_outcome: outcome # type: ignore[method-assign] + runtime = ExecutionRuntime( + lambda gateway=gateway: gateway, action_timeout=0.2, poll_interval=10 + ) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.RUNNING) + runtime._submit( + lambda runtime=runtime: runtime._reduce(_Event("fault", diagnostic="remote fault")) + ) + reset_handle = runtime.reset().value + assert reset_handle is not None + reset = runtime.wait_for_reset(reset_handle, timeout=1) + assert reset.accepted and reset.value is not None and not reset.value.success + assert runtime.snapshot().state == LifecycleState.FAULT + finally: + runtime.shutdown() + + +def test_pending_late_execute_action_makes_reset_fail_permanently() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.01, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + assert gateway.execute_entered["a"].wait(1) + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + reset_handle = runtime.reset().value + assert reset_handle is not None + reset = runtime.wait_for_reset(reset_handle, timeout=1) + assert reset.accepted and reset.value is not None and not reset.value.success + gateway.release["a"].set() + time.sleep(0.05) + assert runtime.snapshot().state == LifecycleState.FAULT + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_failed_reset_can_start_a_new_reconciliation_attempt() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + status_calls = 0 + + def status(_task: str) -> Outcome: + nonlocal status_calls + status_calls += 1 + return Outcome.UNKNOWN if status_calls == 1 else Outcome.INACTIVE + + gateway.status = status # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime._submit(lambda: runtime._reduce(_Event("fault", diagnostic="remote fault"))) + gateway.release["a"].set() + + first = runtime.reset() + assert first.accepted and first.value is not None + first_result = runtime.wait_for_reset(first.value, timeout=1) + assert first_result.accepted and first_result.value is not None + assert not first_result.value.success + assert runtime.snapshot().state == LifecycleState.FAULT + + second = runtime.reset() + assert second.accepted and second.value is not None + assert second.value != first.value + second_result = runtime.wait_for_reset(second.value, timeout=1) + assert second_result.accepted and second_result.value is not None + assert second_result.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_stale_reset_completion_cannot_clear_fault_for_retry() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + first_reset_entered = threading.Event() + release_first_reset = threading.Event() + second_reset_entered = threading.Event() + release_second_reset = threading.Event() + reset_calls = 0 + + def status(_task: str) -> Outcome: + return Outcome.INACTIVE + + def reset_remote(_task: str) -> Outcome: + nonlocal reset_calls + reset_calls += 1 + if reset_calls == 1: + first_reset_entered.set() + assert release_first_reset.wait(2) + elif reset_calls == 2: + second_reset_entered.set() + assert release_second_reset.wait(2) + return Outcome.INACTIVE + + gateway.status = status # type: ignore[method-assign] + gateway.reset = reset_remote # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + operation = runtime.snapshot().operation + assert operation is not None + task = operation.tasks[0] + faulted = replace( + operation, + tasks=(replace(task, activity=TaskActivity.INACTIVE, reset_required=True),), + ) + runtime._submit( + lambda: runtime._commit( + replace( + runtime._context, + state=LifecycleState.FAULT, + fault="remote fault", + diagnostic="remote fault", + active=faulted, + ) + ) + ) + + first = runtime.reset() + assert first.accepted and first.value is not None + assert first_reset_entered.wait(1) + first_result = runtime.wait_for_reset(first.value, timeout=1) + assert first_result.accepted and first_result.value is not None + assert not first_result.value.success + + second = runtime.reset() + assert second.accepted and second.value is not None and second.value != first.value + release_first_reset.set() + wait_until(lambda: reset_calls == 2) + assert second_reset_entered.wait(1) + assert runtime.snapshot().state == LifecycleState.FAULT + assert runtime._context.reset_handle == second.value + + release_second_reset.set() + second_result = runtime.wait_for_reset(second.value, timeout=1) + assert second_result.accepted and second_result.value is not None + assert second_result.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + finally: + release_first_reset.set() + release_second_reset.set() + runtime.shutdown() + + +def test_stale_reset_cancel_cannot_affect_retry_reconciliation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + first_cancel_entered = threading.Event() + release_first_cancel = threading.Event() + second_cancel_entered = threading.Event() + release_second_cancel = threading.Event() + cancel_calls = 0 + + def cancel(_task: str) -> Outcome: + nonlocal cancel_calls + cancel_calls += 1 + if cancel_calls == 1: + first_cancel_entered.set() + assert release_first_cancel.wait(2) + return Outcome.UNKNOWN + second_cancel_entered.set() + assert release_second_cancel.wait(2) + return Outcome.CANCELLED + + gateway.cancel = cancel # type: ignore[method-assign] + gateway.status = lambda _task: Outcome.INACTIVE # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime._submit(lambda: runtime._reduce(_Event("fault", diagnostic="remote fault"))) + + first = runtime.reset() + assert first.accepted and first.value is not None + assert first_cancel_entered.wait(1) + first_result = runtime.wait_for_reset(first.value, timeout=1) + assert first_result.accepted and first_result.value is not None + assert not first_result.value.success + + second = runtime.reset() + assert second.accepted and second.value is not None and second.value != first.value + release_first_cancel.set() + wait_until(lambda: cancel_calls == 2) + assert second_cancel_entered.wait(1) + assert runtime.snapshot().state == LifecycleState.FAULT + assert runtime._context.reset_handle == second.value + + release_second_cancel.set() + second_result = runtime.wait_for_reset(second.value, timeout=1) + assert second_result.accepted and second_result.value is not None + assert second_result.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + finally: + release_first_cancel.set() + release_second_cancel.set() + runtime.shutdown() + + +def test_reset_completion_after_failure_is_retired_before_retry_starts() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + entered = threading.Event() + release = threading.Event() + reset_calls = 0 + + def reset_remote(_task: str) -> Outcome: + nonlocal reset_calls + reset_calls += 1 + if reset_calls == 1: + entered.set() + assert release.wait(2) + return Outcome.INACTIVE + + gateway.reset = reset_remote # type: ignore[method-assign] + gateway.status = lambda _task: Outcome.INACTIVE # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + operation = runtime.snapshot().operation + assert operation is not None + task = operation.tasks[0] + faulted = replace( + operation, + tasks=(replace(task, activity=TaskActivity.INACTIVE, reset_required=True),), + ) + runtime._submit( + lambda: runtime._commit( + replace(runtime._context, state=LifecycleState.FAULT, active=faulted) + ) + ) + first = runtime.reset() + assert first.accepted and first.value is not None and entered.wait(1) + first_result = runtime.wait_for_reset(first.value, timeout=1) + assert first_result.accepted and first_result.value is not None + assert not first_result.value.success + release.set() + wait_until( + lambda: runtime._context.active is not None + and runtime._context.active.tasks[0].action is None + ) + assert runtime.snapshot().state == LifecycleState.FAULT + + second = runtime.reset() + assert second.accepted and second.value is not None and second.value != first.value + wait_until(lambda: reset_calls == 2) + second_result = runtime.wait_for_reset(second.value, timeout=1) + assert second_result.accepted and second_result.value is not None + assert second_result.value.success + finally: + release.set() + runtime.shutdown() + + +def test_stale_reset_deadline_does_not_retire_inflight_action() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + entered = threading.Event() + release = threading.Event() + second_entered = threading.Event() + cancel_calls = 0 + + def cancel(_task: str) -> Outcome: + nonlocal cancel_calls + cancel_calls += 1 + if cancel_calls == 2: + second_entered.set() + entered.set() + assert release.wait(2) + return Outcome.CANCELLED + + gateway.cancel = cancel # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + operation = runtime.snapshot().operation + assert operation is not None + task = operation.tasks[0] + faulted = replace(operation, tasks=(replace(task, activity=TaskActivity.ACTIVE),)) + runtime._submit( + lambda: runtime._commit( + replace(runtime._context, state=LifecycleState.FAULT, active=faulted) + ) + ) + first = runtime.reset() + assert first.accepted and first.value is not None and entered.wait(1) + old_action = runtime._context.active.tasks[0].action + assert old_action is not None and old_action.reset_id == first.value.reset_id + first_result = runtime.wait_for_reset(first.value, timeout=1) + assert first_result.accepted and first_result.value is not None + assert not first_result.value.success + + second = runtime.reset() + assert second.accepted and second.value is not None + runtime._submit( + lambda: runtime._reduce( + _Event( + "action_deadline", + runtime._context.active.handle.operation_id, + runtime._context.active.handle.attempt_id, + task.task_id, + old_action.action_id, + ActionMethod.CANCEL, + reset_id=first.value.reset_id, + ) + ) + ) + assert runtime._context.active.tasks[0].action is old_action + assert runtime._context.reset_handle == second.value + assert cancel_calls == 1 + assert not second_entered.is_set() + release.set() + finally: + release.set() + runtime.shutdown() + + +def test_repeated_reset_calls_share_one_serialized_result() -> None: + gateway = BlockingGateway({"a": Outcome.UNKNOWN}) + entered, release = threading.Event(), threading.Event() + + def status(task: str) -> Outcome: + entered.set() + assert release.wait(2) + return Outcome.UNKNOWN + + gateway.status = status # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + gateway.release["a"].set() + first = runtime.reset() + assert first.accepted and first.value is not None and entered.wait(1) + second = runtime.reset() + assert second.accepted and second.value == first.value + release.set() + result = runtime.wait_for_reset(first.value, timeout=1) + assert result.accepted and result.value is not None and not result.value.success + assert len(runtime._context.reset_results) == 1 + finally: + release.set() + gateway.release["a"].set() + runtime.shutdown() + + +def test_idle_shutdown_is_idempotent_and_rejects_new_work() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + first = runtime.shutdown(timeout=1) + assert first.accepted and first.value is not None and first.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + assert not getattr(runtime.start_planning(), "accepted", False) + assert not runtime.execute_explicit(plan("a")).accepted + second = runtime.shutdown(timeout=1) + assert second.accepted and second.value == first.value + assert gateway.stop_calls == 1 + + +def test_gateway_stop_failure_never_reports_closed_success() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + + def fail_stop() -> None: + gateway.stop_calls += 1 + raise RuntimeError("close exploded") + + gateway.stop = fail_stop # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None and not result.value.success + assert "gateway close failed: close exploded" in result.value.diagnostic + assert result.snapshot is not None + assert result.snapshot.shutdown == ShutdownState.CLOSING + assert result.snapshot.shutdown_result == result.value + assert gateway.stop_calls == 1 + assert runtime.shutdown(timeout=1).accepted is False + assert gateway.stop_calls == 1 + + +def test_shutdown_interrupts_early_action_deadline_thread() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + gateway.release["a"].set() + assert runtime.shutdown(timeout=1).accepted + assert runtime._auxiliary == {} + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_shutdown_interrupts_reset_deadline_thread() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) + handle = ResetHandle("reset-deadline-test") + try: + runtime._start_auxiliary( + runtime._reset_deadline, + handle, + time.monotonic() + 10.0, + name="reset-deadline-test", + ) + wait_until(lambda: len(runtime._auxiliary) == 1) + result = runtime.shutdown(timeout=1) + assert result.accepted + assert runtime._auxiliary == {} + finally: + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_cancel_if_current_cannot_cancel_replacement_operation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + gateway.status = lambda task_name: Outcome.COMPLETED # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + original = runtime.execute_explicit(plan("a")) + assert original.accepted and original.value is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + assert runtime.wait_for_terminal(original.value, timeout=1).accepted + + replacement = runtime.execute_explicit(plan("a")) + assert replacement.accepted and replacement.value is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + + stale_cancel = runtime.cancel_if_current(original.value) + assert not stale_cancel.accepted + assert stale_cancel.diagnostic == "operation is no longer current" + active = runtime.snapshot().operation + assert active is not None + assert active.handle == replacement.value + assert gateway.count("a", ActionMethod.CANCEL) == 0 + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_shutdown_during_pending_cancel_reconciles_and_closes_once() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + assert runtime.cancel().accepted + assert gateway.cancel_entered["a"].wait(1) + result_holder: list[Any] = [] + thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) + thread.start() + wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) + assert not runtime.execute_explicit(plan("a")).accepted + assert not getattr(runtime.start_planning(), "accepted", False) + gateway.release["a"].set() + thread.join(2) + assert ( + len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success + ) + assert gateway.stop_calls == 1 + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown() + + +def test_shutdown_late_execute_acceptance_is_compensated_before_close() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + assert gateway.execute_entered["a"].wait(1) + result_holder: list[Any] = [] + thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) + thread.start() + wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) + gateway.release["a"].set() + thread.join(2) + assert ( + len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success + ) + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown() + + +def test_zero_timeout_does_not_publish_closed_while_owner_is_alive() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + entered, release = threading.Event(), threading.Event() + runtime._events.put(_Call(lambda: (entered.set(), release.wait(1)))) + try: + assert entered.wait(1) + result = runtime.shutdown(timeout=0) + assert not result.accepted and result.value is not None and not result.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + assert runtime._owner.is_alive() + assert not any( + method in (ActionMethod.EXECUTE, ActionMethod.CANCEL) for _, method in gateway.calls + ) + finally: + release.set() + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_zero_timeout_active_shutdown_never_starts_rpc_after_gateway_stop() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + result = runtime.shutdown(timeout=0) + assert not result.accepted and result.value is not None and not result.value.success + time.sleep(0.05) + assert gateway.stop_calls == 1 + assert len(gateway.calls) == gateway.calls_at_stop + assert gateway.count("a", ActionMethod.CANCEL) == 0 + finally: + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_gateway_lease_precedes_closure_and_closer_waits_for_inflight_call() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result_holder: list[tuple[bool, Outcome | None]] = [] + try: + thread = threading.Thread( + target=lambda: result_holder.append( + runtime._invoke_gateway_action("action", "a", ActionMethod.EXECUTE, {"x": 1}) + ) + ) + thread.start() + assert gateway.execute_entered["a"].wait(1) + runtime._seal_gateway() + assert gateway.stop_calls == 0 + gateway.release["a"].set() + thread.join(1) + wait_until(lambda: gateway.stop_calls == 1) + assert result_holder == [(True, Outcome.ACCEPTED)] + assert gateway.calls_at_stop == len(gateway.calls) + finally: + gateway.release["a"].set() + runtime.shutdown(timeout=1) + + +def test_sealed_gateway_rejects_prepared_action_without_inflight_or_rpc() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime._seal_gateway() + admitted, outcome = runtime._invoke_gateway_action( + "action", "a", ActionMethod.EXECUTE, {"x": 1} + ) + assert not admitted and outcome is None + assert gateway.calls == [] + assert runtime._rpc_inflight == 0 + finally: + runtime.shutdown(timeout=1) + + +def test_shutdown_transition_publishes_closing_and_cancel_only_atomically() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + runtime._gateway_condition.acquire() + result_holder: list[Any] = [] + thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) + thread.start() + try: + time.sleep(0.02) + assert thread.is_alive() + assert runtime._context.shutdown == ShutdownState.OPEN + assert runtime._gateway_gate == _GatewayGateState.OPEN + finally: + runtime._gateway_condition.release() + thread.join(2) + assert len(result_holder) == 1 and result_holder[0].accepted + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + + +def test_shutdown_timeout_reports_unresolved_action_and_never_schedules_late_effects() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + assert gateway.execute_entered["a"].wait(1) + result = runtime.shutdown(timeout=0.05) + assert not result.accepted and result.value is not None and not result.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + cancel_count = gateway.count("a", ActionMethod.CANCEL) + gateway.release["a"].set() + time.sleep(0.05) + assert gateway.count("a", ActionMethod.CANCEL) == cancel_count + assert gateway.stop_calls == 1 + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown() + + +def _prep_robot( + name: str, + joints: tuple[str, ...], + group: str, + task: str | None, + mapping: dict[str, str] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path(f"/{name}.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), + joint_names=list(joints), + planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], + coordinator_task_name=task, + joint_name_mapping=mapping or {}, + ) + + +def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: + return GeneratedPlan( + group_ids=groups, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width + ), + TrajectoryPoint( + time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width + ), + ], + ), + ) + + +def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + generated = _generated( + ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 + ) + + prepared = prepare_generated_plan(generated, topology) + assert prepared.generated_plan is generated + assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] + assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] + assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] + assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 + assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] + + +def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + topology = ExecutionTopology.from_robot_configs((left, right)) + with pytest.raises(ValueError, match="route set"): + ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) + with pytest.raises(ValueError, match="malformed"): + prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) + with pytest.raises(ValueError, match="missing robot joints"): + prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) + with pytest.raises(ValueError, match="duplicate global"): + prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + + +def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + assert isinstance(prepared, PreparedPlan) + assert prepared is not generated + partial = PreparedPlan(generated, (), topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + assert not runtime.execute_explicit(partial).accepted + finally: + runtime.shutdown() + + +def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + legacy = ExecutionPlan(generated, prepared.entries) + assert not runtime.execute_explicit(legacy).accepted + assert runtime.execute_explicit(prepared).accepted + finally: + runtime.shutdown() + + +def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: + left = _prep_robot("left", ("l0",), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + prepared = prepare_generated_plan( + _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology + ) + runtime = ExecutionRuntime( + lambda: BlockingGateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), + poll_interval=10, + ) + try: + result = runtime.execute_explicit(prepared) + assert result.accepted + assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] + finally: + runtime.shutdown() + + +def test_prepare_rejects_additional_selected_robot_joint_column() -> None: + robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + with pytest.raises(ValueError, match="additional joint column"): + prepare_generated_plan( + _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology + ) + + +def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + with pytest.raises(ValueError, match="duplicate group route"): + ExecutionTopology.from_robot_configs( + (robot,), + [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], + ) + with pytest.raises(ValueError, match="mapped local joint"): + ExecutionTopology.from_robot_configs( + (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) + ) + bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + bad_group.planning_groups[0] = PlanningGroupDefinition( + name="manipulator", joint_names=("missing",), base_link="base" + ) + with pytest.raises(ValueError, match="planning group joints"): + ExecutionTopology.from_robot_configs((bad_group,)) + + +def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: + planning_only = _prep_robot("planner", ("j0",), "manipulator", None) + topology = ExecutionTopology.from_robot_configs((planning_only,)) + assert topology.routes == () + with pytest.raises(ValueError, match="no coordinator task"): + prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) + + +def test_clear_ready_plan_is_a_dedicated_ready_to_idle_transition() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + token = runtime.start_planning() + assert token is not None + assert runtime.complete_planning(token, plan("a")).accepted + assert runtime.snapshot().state == LifecycleState.READY + cleared = runtime.clear_ready_plan() + assert cleared.accepted + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.ready_plan is None and snapshot.ready_plan_id is None + assert not runtime.clear_ready_plan().accepted + finally: + runtime.shutdown() + + +def test_runtime_diagnostic_persists_idle_planning_and_execution_errors_then_clears_on_planning() -> ( + None +): + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) + try: + token = runtime.start_planning() + assert token is not None + invalid = runtime.complete_planning(token, ExecutionPlan("generated", ())) + assert not invalid.accepted + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == "invalid prepared plan" + assert runtime.start_planning() is not None + assert runtime.snapshot().diagnostic is None + finally: + runtime.shutdown() + + +def test_rejected_execute_retains_contextual_diagnostic_in_snapshot_and_result() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a")).value + assert handle is not None + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.REJECTED + assert "task=a" in result.value.diagnostic + assert "action=" in result.value.diagnostic + assert "outcome=rejected" in result.value.diagnostic + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.diagnostic == result.value.diagnostic + finally: + runtime.shutdown() + + +def test_status_terminal_failure_retains_contextual_diagnostic_in_terminal_result() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) + gateway.status = lambda task: Outcome.INACTIVE if task == "a" else Outcome.RUNNING # type: ignore[method-assign] + gateway.release["a"].set() + gateway.release["b"].set() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + runtime.poll() + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.FAILED + assert "task=a" in result.value.diagnostic + assert "outcome=inactive" in result.value.diagnostic + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == result.value.diagnostic + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_auxiliary_gripper_calls_use_private_gateway_and_normalize_results() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + set_result = runtime.set_gripper_position("gripper", 0.75) + get_result = runtime.get_gripper_position("gripper") + assert set_result.accepted and set_result.value == Outcome.ACCEPTED + assert get_result.accepted and get_result.value == 0.25 + assert gateway.count("gripper", ActionMethod.GRIPPER_SET) == 1 + assert gateway.count("gripper", ActionMethod.GRIPPER_GET) == 1 + assert not hasattr(runtime, "gateway") + finally: + runtime.shutdown() + + +def test_control_gateway_uses_top_level_gripper_rpc_methods() -> None: + rpc = FakeRPC( + {("gripper", "set_gripper_position"): True, ("gripper", "get_gripper_position"): 0.6}, [] + ) + gateway = ControlCoordinatorGateway(rpc) + assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED + assert gateway.get_gripper_position("gripper") == 0.6 + assert rpc.calls == [ + ("gripper", "set_gripper_position", {"position": 0.4}), + ("gripper", "get_gripper_position", {}), + ] + + +def test_auxiliary_calls_reject_after_closing_and_inflight_auxiliary_drains_before_close() -> None: + gateway = BlockingAuxGateway() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + results: list[Any] = [] + caller = threading.Thread( + target=lambda: results.append(runtime.set_gripper_position("gripper", 0.5)) + ) + caller.start() + assert gateway.aux_entered.wait(1) + shutdown_results: list[Any] = [] + closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(timeout=1))) + closer.start() + wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) + gateway.aux_release.set() + caller.join(2) + closer.join(2) + assert results and results[0].accepted + assert shutdown_results and shutdown_results[0].accepted + assert gateway.calls_at_stop == len(gateway.calls) + assert gateway.stop_calls == 1 + assert not runtime.set_gripper_position("gripper", 0.1).accepted + assert not runtime.get_gripper_position("gripper").accepted + + +def test_gateway_factory_is_the_internal_runtime_injection_seam() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) + factories: list[bool] = [] + + def factory() -> BlockingGateway: + factories.append(True) + return gateway + + runtime = ExecutionRuntime(factory, poll_interval=10) + try: + assert factories == [True] + assert runtime.get_gripper_position("gripper").accepted + assert not hasattr(runtime, "client") + finally: + runtime.shutdown() diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py index 01d52ae89b..dd341d8cc0 100644 --- a/dimos/manipulation/test_generated_plan_materialization.py +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -14,13 +14,18 @@ """Generated-plan materialization tests.""" +from collections.abc import Iterator from pathlib import Path from unittest.mock import MagicMock import pytest -from dimos.manipulation._test_manipulation_helpers import make_module -from dimos.manipulation.manipulation_module import ManipulationState +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module, +) +from dimos.manipulation.execution_runtime import PreparedPlan from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig @@ -75,9 +80,16 @@ def _robot(name: str, joints: list[str], velocity: float, acceleration: float) - ], max_velocity=velocity, max_acceleration=acceleration, + coordinator_task_name=f"task_{name}", ) +@pytest.fixture(autouse=True) +def _close_runtime() -> Iterator[None]: + yield + close_test_runtimes() + + def _module(monkeypatch: pytest.MonkeyPatch): RecordingGenerator.calls = [] RecordingGenerator.limits = None @@ -96,8 +108,7 @@ def _module(monkeypatch: pytest.MonkeyPatch): module._world_monitor.world = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) module._planner = MagicMock() - module._state = ManipulationState.PLANNING - module._planning_epoch = 1 + install_runtime(module, [left, right]) return module @@ -115,13 +126,16 @@ def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distin status=PlanningStatus.SUCCESS, path=path ) - assert module._plan_selected_path(("left/group", "right/group"), path[0], path[-1], 1) + token = module._begin_group_planning() + assert token is not None + assert module._plan_selected_path(("left/group", "right/group"), path[0], path[-1], token) assert RecordingGenerator.calls == [[[0.0, 0.0, 0.0], [0.2, 0.1, 0.3]]] assert RecordingGenerator.limits == ([1.0, 1.0, 3.0], [2.0, 2.0, 4.0]) - assert module._last_plan is not None - assert module._last_plan.path is not module._last_plan.trajectory.points - assert module._last_plan.trajectory.joint_names == names - assert module._last_plan.trajectory.points[-1].time_from_start == 1.0 + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.generated_plan.path is not ready.generated_plan.trajectory.points + assert ready.generated_plan.trajectory.joint_names == names + assert ready.generated_plan.trajectory.points[-1].time_from_start == 1.0 @pytest.mark.parametrize( @@ -138,9 +152,11 @@ def test_rejects_malformed_or_nonfinite_waypoints(monkeypatch, path, message): status=PlanningStatus.SUCCESS, path=path ) - assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) - assert module._last_plan is None - assert message in module._error_message + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None + assert message in (module.get_error() or "") def test_rejects_invalid_limits_and_generator_failure_without_caching(monkeypatch): @@ -152,8 +168,10 @@ def test_rejects_invalid_limits_and_generator_failure_without_caching(monkeypatc status=PlanningStatus.SUCCESS, path=path ) - assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) - assert module._last_plan is None + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None assert RecordingGenerator.calls == [] module = _module(monkeypatch) @@ -161,8 +179,10 @@ def test_rejects_invalid_limits_and_generator_failure_without_caching(monkeypatc module._planner.plan_selected_joint_path.return_value = PlanningResult( status=PlanningStatus.SUCCESS, path=path ) - assert not module._plan_selected_path(("left/group",), path[0], path[-1], 1) - assert module._last_plan is None + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None assert len(RecordingGenerator.calls) == 1 @@ -173,9 +193,10 @@ def test_zero_generation_after_caching_for_status_and_completion(monkeypatch): module._planner.plan_selected_joint_path.return_value = PlanningResult( status=PlanningStatus.SUCCESS, path=path ) - assert module._plan_selected_path(("left/group",), path[0], path[-1], 1) + token = module._begin_group_planning() + assert token is not None + assert module._plan_selected_path(("left/group",), path[0], path[-1], token) RecordingGenerator.calls = [] - module._get_coordinator_client = MagicMock(return_value=None) - module._wait_for_trajectory_completion("left", timeout=0.0) + assert module.get_trajectory_status() is not None assert RecordingGenerator.calls == [] diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index 4a6eb99c63..b1f49344dd 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -22,14 +22,17 @@ from __future__ import annotations import importlib.util -from unittest.mock import MagicMock import pytest -from dimos.manipulation.manipulation_module import ( - ManipulationModule, - ManipulationState, +from dimos.manipulation._test_manipulation_helpers import FakeCoordinatorGateway +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + LifecycleState, + Outcome, + PreparedPlan, ) +from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.Pose import Pose @@ -122,6 +125,9 @@ def module(xarm7_config): planning_timeout=10.0, visualization={"backend": "none"}, ) + gateway = FakeCoordinatorGateway() + mod._runtime_factory = lambda _factory, **kwargs: ExecutionRuntime(lambda: gateway, **kwargs) + mod._runtime_gateway = gateway mod.coordinator_joint_state = None mod.objects = None mod.start() @@ -136,7 +142,7 @@ class TestManipulationModuleIntegration: def test_module_initialization(self, module): """Test module initializes with real Drake world.""" - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name assert module._world_monitor is not None assert module._planner is not None assert module._kinematics is not None @@ -166,14 +172,14 @@ def test_plan_to_joints(self, module, joint_state_zeros): success = module.plan_to_joints(target) assert success is True - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.READY.name assert module.has_planned_path() is True - assert module._last_plan is not None - traj = module._split_plan_trajectory_by_robot(module._last_plan)["test_arm"] - assert len(traj.points) > 1 - assert traj.duration > 0 - assert module._last_plan.group_ids == ("test_arm/manipulator",) + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert len(ready.generated_plan.trajectory.points) > 1 + assert ready.generated_plan.trajectory.duration > 0 + assert ready.generated_plan.group_ids == ("test_arm/manipulator",) def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): """Test planning to an explicit planning-group joint target.""" @@ -184,11 +190,11 @@ def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): ) assert success is True - assert module._state == ManipulationState.COMPLETED - assert module._last_plan is not None - assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module.get_state() == LifecycleState.READY.name + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.generated_plan.group_ids == ("test_arm/manipulator",) assert module.has_planned_path() is True - assert module._split_plan_trajectory_by_robot(module._last_plan) is not None def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" @@ -234,26 +240,24 @@ def test_ee_pose(self, module, joint_state_zeros): assert hasattr(pose, "y") assert hasattr(pose, "z") - def test_trajectory_name_translation(self, module, joint_state_zeros): - """Test that trajectory joint names are translated for coordinator.""" + def test_prepared_plan_contains_coordinator_joint_names(self, module, joint_state_zeros): + """Runtime materialization translates selected joints for the gateway.""" module._on_joint_state(joint_state_zeros) success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True - assert module._last_plan is not None - traj = module._split_plan_trajectory_by_robot(module._last_plan)["test_arm"] - robot_config = module._robots["test_arm"][1] - - translated = module._translate_trajectory_to_coordinator(traj, robot_config) - - assert translated.joint_names == list(robot_config.joint_name_mapping.keys()) + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.entries[0].request["trajectory"].joint_names == list( + module._robots["test_arm"][1].joint_name_mapping.values() + ) @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") class TestCoordinatorIntegration: - """Test coordinator integration with mocked RPC client.""" + """Test coordinator integration through the runtime gateway.""" def test_execute_with_mock_coordinator(self, module, joint_state_zeros): """Test execute sends trajectory to coordinator.""" @@ -262,28 +266,18 @@ def test_execute_with_mock_coordinator(self, module, joint_state_zeros): success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True - # Mock the coordinator client - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - result = module.execute() assert result is True - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.RUNNING.name - # Verify coordinator was called - mock_client.task_invoke.assert_called_once() - call_args = mock_client.task_invoke.call_args - task_name, method_name, kwargs = call_args[0] + task_name, request = module._runtime_gateway.execute_calls[0] assert task_name == "traj_arm" - assert method_name == "execute" - trajectory = kwargs["trajectory"] + trajectory = request["trajectory"] assert len(trajectory.points) > 1 - # Joint names should be translated robot_config = module._robots["test_arm"][1] - assert trajectory.joint_names == list(robot_config.joint_name_mapping.keys()) + assert trajectory.joint_names == list(robot_config.joint_name_mapping.values()) def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): """Test handling of coordinator rejection.""" @@ -291,39 +285,34 @@ def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): module.plan_to_joints(JointState(position=[0.05] * 7)) - # Mock coordinator to reject - mock_client = MagicMock() - mock_client.task_invoke.return_value = False - module._coordinator_client = mock_client + module._runtime_gateway.execute_outcome = Outcome.REJECTED result = module.execute() assert result is False - assert module._state == ManipulationState.FAULT - assert "rejected" in module._error_message.lower() + assert module.get_state() == LifecycleState.IDLE.name + assert "rejected" in module.get_error().lower() + terminal_results = dict(module._execution_runtime._context.terminal_results) + assert terminal_results + assert next(reversed(terminal_results.values())).outcome == Outcome.REJECTED def test_state_transitions_during_execution(self, module, joint_state_zeros): """Test state transitions during plan and execute.""" - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name module._on_joint_state(joint_state_zeros) - # Plan - should go through PLANNING -> COMPLETED + # Plan commits a runtime-owned READY plan. module.plan_to_joints(JointState(position=[0.05] * 7)) - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.READY.name - # Reset works from COMPLETED + # Reset clears the READY plan. module.reset() - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name # Plan again module.plan_to_joints(JointState(position=[0.05] * 7)) - # Mock coordinator - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - - # Execute - should go to EXECUTING then COMPLETED + # Execute is accepted by the runtime gateway. module.execute() - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.RUNNING.name diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py index 06cea02e45..2e067f14ad 100644 --- a/dimos/manipulation/test_manipulation_monitor_preview.py +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -16,12 +16,18 @@ from __future__ import annotations +from collections.abc import Iterator from pathlib import Path from unittest.mock import MagicMock import pytest -from dimos.manipulation._test_manipulation_helpers import make_module as _make_module +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module as _make_module, +) +from dimos.manipulation.execution_runtime import prepare_generated_plan from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry @@ -38,6 +44,12 @@ from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +@pytest.fixture(autouse=True) +def _close_test_runtimes() -> Iterator[None]: + yield + close_test_runtimes() + + @pytest.fixture def robot_config_with_mapping() -> RobotModelConfig: """Create a robot config with joint name mapping.""" @@ -95,7 +107,7 @@ def _install_generated_plan( name=config.joint_names, position=[0.0 for _ in config.joint_names], ) - module._last_plan = GeneratedPlan( + generated_plan = GeneratedPlan( trajectory=JointTrajectory( joint_names=global_joint_names, points=[ @@ -117,6 +129,11 @@ def _install_generated_plan( for point in points ], ) + install_runtime(module, [config]) + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(generated_plan, module._execution_topology) + assert token is not None + assert module._execution_runtime.complete_planning(token, prepared).accepted def _make_module_with_monitor(*configs: RobotModelConfig) -> ManipulationModule: @@ -196,193 +213,14 @@ def close(self) -> None: class TestOnJointState: """Test _on_joint_state routing, splitting, and init capture.""" - def test_routes_positions_to_monitor(self, robot_config_with_mapping): - """Joint positions from aggregated message are routed to the correct monitor.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - velocity=[1.0, 2.0, 3.0], - ) - module._on_joint_state(msg) - - # Verify world_monitor received the sub-message - module._world_monitor.on_joint_state.assert_called_once() - call_args = module._world_monitor.on_joint_state.call_args - sub_msg = call_args[0][0] - assert sub_msg.position == [0.1, 0.2, 0.3] - assert sub_msg.velocity == [1.0, 2.0, 3.0] - assert call_args[1]["robot_id"] == "robot_left_arm" - - def test_skips_robot_with_missing_joints(self, robot_config_with_mapping): - """Robots whose joints are absent from the message are skipped.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - # Message has none of left_arm's joints - msg = JointState( - name=["right/joint1", "right/joint2"], - position=[0.5, 0.6], - ) - module._on_joint_state(msg) - - module._world_monitor.on_joint_state.assert_not_called() - - def test_captures_init_joints_on_first_call(self, robot_config_with_mapping): - """First joint state is stored as init joints; subsequent calls don't overwrite.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - first_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(first_msg) - assert "left_arm" in module._init_joints - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - # Second call should NOT overwrite - second_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.9, 0.8, 0.7], - ) - module._on_joint_state(second_msg) - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - def test_multi_robot_splits_correctly(self): - """With two robots, each gets only its own joints from the aggregated message.""" - left_config = RobotModelConfig( - name="left", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, - coordinator_task_name="traj_left", - ) - right_config = RobotModelConfig( - name="right", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, - coordinator_task_name="traj_right", - ) - module = _make_module_with_monitor(left_config, right_config) - - msg = JointState( - name=["left/j1", "left/j2", "right/j1", "right/j2"], - position=[1.0, 2.0, 3.0, 4.0], - velocity=[0.1, 0.2, 0.3, 0.4], - ) - module._on_joint_state(msg) - - assert module._world_monitor.on_joint_state.call_count == 2 - - # Collect calls by robot_id - calls = { - call[1]["robot_id"]: call[0][0] - for call in module._world_monitor.on_joint_state.call_args_list - } - assert calls["robot_left"].position == [1.0, 2.0] - assert calls["robot_right"].position == [3.0, 4.0] - assert calls["robot_left"].velocity == [0.1, 0.2] - assert calls["robot_right"].velocity == [0.3, 0.4] - - def test_no_monitor_returns_early(self, robot_config_with_mapping): - """When world_monitor is None, _on_joint_state returns without error.""" - module = _make_module() - module._robots = {"left_arm": ("id", robot_config_with_mapping, MagicMock())} - module._world_monitor = None - - # Should not raise - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(msg) + pass class TestWorldMonitorVisualization: - def test_visualization_routing_and_stop_all_monitors(self): - viz = FakeVisualization() - monitor = _make_world_monitor_with_viz(viz) - state_monitor = MagicMock() - obstacle_monitor = MagicMock() - monitor._state_monitors = {"robot": state_monitor} - monitor._obstacle_monitor = obstacle_monitor - monitor._viz_thread = MagicMock() - monitor._viz_thread.is_alive.return_value = False - - assert monitor.get_visualization_url() == "123" - monitor.update_visualization_state() - monitor.cancel_preview_animation() - path = _make_path([1.0], [2.0], [3.0]) - plan = GeneratedPlan( - trajectory=JointTrajectory(), - group_ids=("robot/group",), - status=PlanningStatus.SUCCESS, - path=path, - ) - monitor.animate_trajectory(plan.trajectory, 4.5) - assert monitor.visualization is viz - assert viz.published is True - assert viz.preview_animation_cancellations == 2 - assert viz.animations == [(tuple(), [], 4.5)] - - monitor.stop_all_monitors() - - assert viz.close_count == 1 - state_monitor.stop.assert_called_once() - obstacle_monitor.stop.assert_called_once() - - def test_visualization_none_is_noop(self): - monitor = _make_world_monitor_with_viz(None) - - assert monitor.get_visualization_url() is None - monitor.update_visualization_state() - monitor.cancel_preview_animation() - monitor.animate_trajectory(JointTrajectory(), 1.0) - monitor.start_visualization_thread() - assert monitor._viz_thread is None + pass class TestManipulationPreview: - def test_clear_planned_path_invalidates_before_dismissing_preview(self): - module = _make_module() - plan = GeneratedPlan(trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[]) - module._last_plan = plan - module._world_monitor = MagicMock() - plan_during_dismissal: list[GeneratedPlan | None] = [] - module._world_monitor.cancel_preview_animation.side_effect = ( - lambda: plan_during_dismissal.append(module._last_plan) - ) - - assert module.clear_planned_path() is True - - assert plan_during_dismissal == [None] - module._world_monitor.cancel_preview_animation.assert_called_once_with() - assert module._last_plan is None - - def test_clear_planned_path_clears_without_a_world_monitor(self): - module = _make_module() - module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] - ) - - assert module.clear_planned_path() is True - assert module._last_plan is None - def test_dismiss_preview_noop_without_monitor(self): module = _make_module() @@ -396,45 +234,6 @@ def test_dismiss_preview_routes_to_monitor(self): module._world_monitor.cancel_preview_animation.assert_called_once_with() - def test_preview_routes_one_complete_plan_with_default_duration(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - _install_generated_plan(module, config, traj_gen, [0.0], [2.0]) - - assert module.preview_plan() is True - - module._world_monitor.animate_trajectory.assert_called_once_with( - module._last_plan.trajectory, None - ) - - def test_preview_robot_name_validates_affectedness_without_trimming(self): - module = _make_module() - left = _one_joint_config("left") - right = _one_joint_config("right") - traj_gen = MagicMock() - module._robots = { - "left": ("left_id", left, traj_gen), - "right": ("right_id", right, traj_gen), - } - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) - module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), - group_ids=("left/manipulator", "right/manipulator"), - status=PlanningStatus.SUCCESS, - path=[ - JointState(name=["left/j0", "right/j0"], position=[0.0, 0.0]), - JointState(name=["left/j0", "right/j0"], position=[1.0, 1.0]), - ], - ) - - assert module.preview_plan(duration=2.5, robot_name="left") is True - - module._world_monitor.animate_trajectory.assert_called_once_with( - module._last_plan.trajectory, 2.5 - ) - def test_preview_rejects_unaffected_compatibility_robot(self): module = _make_module() config = _one_joint_config() diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 4ad998517e..fc00666169 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -16,17 +16,22 @@ from __future__ import annotations +from collections.abc import Iterator from pathlib import Path from unittest.mock import MagicMock import pytest from pytest_mock import MockerFixture -from dimos.manipulation._test_manipulation_helpers import make_module as _make_module +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module as _make_module, +) +from dimos.manipulation.execution_runtime import prepare_generated_plan from dimos.manipulation.manipulation_module import ( ManipulationModule, ManipulationModuleConfig, - ManipulationState, ) from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry @@ -36,8 +41,6 @@ from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus from dimos.manipulation.planning.spec.models import ( GeneratedPlan, - IKResult, - PlanningResult, ) from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -71,6 +74,12 @@ def robot_config(): ) +@pytest.fixture(autouse=True) +def _close_test_runtimes() -> Iterator[None]: + yield + close_test_runtimes() + + @pytest.fixture def robot_config_with_mapping(): """Create a robot config with joint name mapping (dual-arm scenario).""" @@ -148,7 +157,7 @@ def _install_generated_plan( name=global_joint_names, position=[0.0 for _ in config.joint_names], ) - module._last_plan = GeneratedPlan( + generated_plan = GeneratedPlan( trajectory=JointTrajectory( joint_names=global_joint_names, points=[ @@ -170,6 +179,11 @@ def _install_generated_plan( for point in points ], ) + install_runtime(module, [config]) + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(generated_plan, module._execution_topology) + assert token is not None + assert module._execution_runtime.complete_planning(token, prepared).accepted def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> JointTrajectory: @@ -200,93 +214,6 @@ def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: class TestStateMachine: """Test state transitions.""" - def test_cancel_interrupts_active_work(self): - """Cancel works for executing motion and in-progress planning.""" - module = _make_module() - - module._state = ManipulationState.IDLE - assert module.cancel() is False - - module._state = ManipulationState.PLANNING - assert module.cancel() is True - assert module._state == ManipulationState.IDLE - assert module._planning_epoch == 1 - - module._state = ManipulationState.EXECUTING - assert module.cancel() is True - assert module._state == ManipulationState.IDLE - - def test_cancel_hides_active_plan_preview(self): - module = _make_module() - module._state = ManipulationState.EXECUTING - module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] - ) - module._world_monitor = MagicMock() - - assert module.cancel() is True - - module._world_monitor.cancel_preview_animation.assert_called_once_with() - - def test_cancel_completed_execution_cancels_coordinator_task(self): - module = _make_module() - module._state = ManipulationState.COMPLETED - module._possibly_active_tasks = {"traj_arm"} - module._robots = { - "arm": ("arm_id", MagicMock(coordinator_task_name="traj_arm"), MagicMock()) - } - module._coordinator_client = MagicMock() - module._coordinator_client.task_invoke.return_value = False - - assert module.cancel() is True - module._coordinator_client.task_invoke.assert_called_once_with("traj_arm", "cancel", {}) - assert module._state == ManipulationState.IDLE - - def test_reset_not_during_execution(self): - """Reset works in any state except EXECUTING.""" - module = _make_module() - - module._state = ManipulationState.FAULT - module._error_message = "Error" - result = module.reset() - assert result.is_success() - assert module._state == ManipulationState.IDLE - assert module._error_message == "" - - module._state = ManipulationState.EXECUTING - result = module.reset() - assert not result.is_success() - assert result.error_code == "INVALID_STATE" - - def test_fail_sets_fault_state(self): - """_fail helper sets FAULT state and message.""" - module = _make_module() - module._state = ManipulationState.PLANNING - - result = module._fail("Test error") - assert result is False - assert module._state == ManipulationState.FAULT - assert module._error_message == "Test error" - - def test_begin_planning_state_checks(self, robot_config): - """_begin_planning only allowed from IDLE or COMPLETED.""" - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} - - # From IDLE - OK - module._state = ManipulationState.IDLE - assert module._begin_planning() == ("test_arm", "robot_id") - assert module._state == ManipulationState.PLANNING - - # From COMPLETED - OK - module._state = ManipulationState.COMPLETED - assert module._begin_planning() == ("test_arm", "robot_id") - - # From EXECUTING - Fail - module._state = ManipulationState.EXECUTING - assert module._begin_planning() is None - class TestRobotSelection: """Test robot selection logic.""" @@ -373,6 +300,7 @@ def test_kinematics_config_is_passed_to_factory( kinematics_name=None, kinematics=kinematics, ) + module.stop() def test_legacy_kinematics_name_still_selects_backend( self, robot_config, planning_initialization: PlanningInitializationHarness @@ -393,6 +321,7 @@ def test_legacy_kinematics_name_still_selects_backend( kinematics_name="pink", kinematics=module.config.kinematics, ) + module.stop() def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: """Pydantic parses the nested CLI config shape used by -o overrides.""" @@ -410,86 +339,6 @@ def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: assert config.kinematics.dt == 0.02 assert config.kinematics.posture_cost == 0.0 - def test_solve_ik_rpc_calls_configured_backend(self, robot_config): - """solve_ik returns the backend IKResult without path planning.""" - 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]) - current = JointState(name=robot_config.joint_names, position=[0.0, 0.0, 0.0]) - current_global = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = current_global - expected = IKResult( - status=IKStatus.SUCCESS, - joint_state=JointState(name=robot_config.joint_names, position=[0.1, 0.2, 0.3]), - position_error=0.0001, - orientation_error=0.0002, - iterations=3, - message="ok", - ) - module._kinematics = MagicMock() - module._kinematics.solve_pose_targets.return_value = expected - - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose) - - assert result is expected - assert module._state == ManipulationState.COMPLETED - assert module._last_plan is None - module._kinematics.solve_pose_targets.assert_called_once() - _, kwargs = module._kinematics.solve_pose_targets.call_args - assert kwargs["world"] is module._world_monitor.world - assert kwargs["seed"].name == current_global.name - assert kwargs["seed"].position == current.position - assert kwargs["check_collision"] is True - [(group, target_pose)] = kwargs["pose_targets"].items() - assert group.id == "test_arm/manipulator" - assert target_pose.frame_id == "world" - assert target_pose.position.x == 0.45 - - def test_solve_ik_rpc_returns_failure_without_joint_state(self, robot_config): - """solve_ik reports a failed IKResult when no seed state is available.""" - module = _make_module() - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=[], position=[] - ) - module._kinematics = MagicMock() - - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose) - - assert result.status == IKStatus.NO_SOLUTION - assert result.message == "No joint state" - assert module._state == ManipulationState.IDLE - module._kinematics.solve_pose_targets.assert_not_called() - - def test_solve_ik_rpc_uses_explicit_seed(self, robot_config): - """solve_ik initializes the backend from an explicit seed when provided.""" - 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]) - explicit_seed = JointState(name=robot_config.joint_names, position=[0.2, 0.1, 0.0]) - expected = IKResult(status=IKStatus.SUCCESS, joint_state=explicit_seed) - module._kinematics = MagicMock() - module._kinematics.solve_pose_targets.return_value = expected - - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose, seed=explicit_seed) - - assert result is expected - _, kwargs = module._kinematics.solve_pose_targets.call_args - assert kwargs["seed"] is explicit_seed - module._world_monitor.current_global_joint_state.assert_not_called() - class TestPlanningGroupApis: """Test explicit planning-group API behavior.""" @@ -511,412 +360,6 @@ def test_list_planning_groups_and_robot_info_include_groups(self, robot_config): assert info["end_effector_link"] == "link_tcp" assert info["has_joint_name_mapping"] is False - def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches(self, robot_config): - module = _make_module() - registry = PlanningGroupRegistry([robot_config]) - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory( - (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) - ) - module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} - module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - result_path = [ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ), - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], - ), - ] - module._planner = MagicMock() - module._planner.plan_selected_joint_path.return_value = PlanningResult( - status=PlanningStatus.SUCCESS, - path=result_path, - planning_time=0.01, - path_length=0.3, - iterations=4, - message="ok", - ) - - success = module.plan_to_joint_targets( - { - "test_arm/manipulator": JointState( - name=robot_config.joint_names, - position=[0.1, 0.2, 0.3], - ) - } - ) - - assert success is True - assert module._last_plan is not None - assert module._last_plan.group_ids == ("test_arm/manipulator",) - assert module._last_plan.path == result_path - trajectories = module._split_plan_trajectory_by_robot(module._last_plan) - assert trajectories is not None - assert trajectories["test_arm"].points[-1].positions == [0.1, 0.2, 0.3] - module._planner.plan_selected_joint_path.assert_called_once() - _, kwargs = module._planner.plan_selected_joint_path.call_args - assert kwargs["selection"].group_ids == ("test_arm/manipulator",) - assert kwargs["goal"].name == [ - "test_arm/joint1", - "test_arm/joint2", - "test_arm/joint3", - ] - - success = module.plan_to_joint_targets( - { - "test_arm/manipulator": JointState( - name=robot_config.joint_names, - position=[0.1, 0.2, 0.3], - ) - } - ) - - assert success is True - assert module._planner.plan_selected_joint_path.call_count == 2 - - def test_plan_to_pose_targets_uses_group_ik_and_selected_path(self, robot_config): - module = _make_module() - registry = PlanningGroupRegistry([robot_config]) - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory( - (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) - ) - module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} - module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - ik_goal = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._kinematics = MagicMock() - module._kinematics.solve_pose_targets.return_value = IKResult( - status=IKStatus.SUCCESS, - joint_state=ik_goal, - ) - module._planner = MagicMock() - module._planner.plan_selected_joint_path.return_value = PlanningResult( - status=PlanningStatus.SUCCESS, - path=[ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ), - ik_goal, - ], - ) - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - - success = module.plan_to_pose_targets({"test_arm/manipulator": pose}) - - assert success is True - module._kinematics.solve_pose_targets.assert_called_once() - _, ik_kwargs = module._kinematics.solve_pose_targets.call_args - target_groups = list(ik_kwargs["pose_targets"].keys()) - assert [group.id for group in target_groups] == ["test_arm/manipulator"] - target_pose = ik_kwargs["pose_targets"][target_groups[0]] - assert target_pose.position.x == 0.45 - assert ik_kwargs["seed"].name == [ - "test_arm/joint1", - "test_arm/joint2", - "test_arm/joint3", - ] - _, planner_kwargs = module._planner.plan_selected_joint_path.call_args - assert planner_kwargs["goal"] is ik_goal - - def test_failed_plan_materialization_clears_generated_plan( - self, robot_config, simple_trajectory - ): - module = _make_module() - registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} - module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = None - module._last_plan = GeneratedPlan( - trajectory=_generated_plan_trajectory( - ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - [0.0, 0.0, 0.0], - [0.1, 0.2, 0.3], - ), - group_ids=("test_arm/manipulator",), - status=PlanningStatus.SUCCESS, - path=[ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ), - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.2, 0.2, 0.2], - ), - ], - ) - module._planner = MagicMock() - module._planner.plan_selected_joint_path.return_value = PlanningResult( - status=PlanningStatus.SUCCESS, - path=[ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], - ) - ], - ) - - success = module.plan_to_joint_targets( - {"test_arm/manipulator": JointState(position=[0.1, 0.2, 0.3])} - ) - - assert success is False - assert module._state == ManipulationState.FAULT - assert module._last_plan is None - assert module.has_planned_path() is False - - def test_execute_plan_dispatches_each_affected_robot(self, robot_config, simple_trajectory): - module = _make_module() - registry = PlanningGroupRegistry([robot_config]) - traj_gen = MagicMock() - traj_gen.generate.return_value = simple_trajectory - module._robots = {"test_arm": ("robot_id", robot_config, traj_gen)} - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = registry - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._last_plan = GeneratedPlan( - trajectory=_generated_plan_trajectory( - ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - [0.0, 0.0, 0.0], - [0.1, 0.2, 0.3], - ), - group_ids=("test_arm/manipulator",), - status=PlanningStatus.SUCCESS, - path=[ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ), - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], - ), - ], - ) - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - - assert module.execute_plan() is True - - mock_client.task_invoke.assert_called_once() - task_name, method_name, payload = mock_client.task_invoke.call_args.args - assert task_name == "traj_arm" - assert method_name == "execute" - trajectory = payload["trajectory"] - assert trajectory.joint_names == simple_trajectory.joint_names - assert [point.positions for point in trajectory.points] == [ - [0.0, 0.0, 0.0], - [0.1, 0.2, 0.3], - ] - traj_gen.generate.assert_not_called() - assert module._state == ManipulationState.COMPLETED - - def test_execute_plan_rejects_stale_planned_robot_without_consuming_plan(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) - module._state = ManipulationState.COMPLETED - module._world_monitor.current_global_joint_state.return_value = JointState( - name=[], position=[] - ) - module._world_monitor.is_state_stale.return_value = True - mock_client = MagicMock() - module._coordinator_client = mock_client - - assert module.execute_plan() is False - - assert module._last_plan is not None - assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_not_called() - - def test_execute_plan_rejects_duplicate_current_global_joints_without_consuming_plan(self): - module = _make_module() - config = _one_joint_config() - traj_gen = MagicMock() - _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) - module._state = ManipulationState.COMPLETED - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["arm/j0", "arm/j0"], position=[0.0, 0.0] - ) - mock_client = MagicMock() - module._coordinator_client = mock_client - - assert module.execute_plan() is False - - assert module._last_plan is not None - assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_not_called() - - def test_execute_plan_rejects_reordered_current_global_joints_without_consuming_plan(self): - module = _make_module() - config = RobotModelConfig( - name="arm", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j0", "j1"], - base_link="base_link", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", - joint_names=("j0", "j1"), - base_link="base_link", - tip_link="ee", - ) - ], - coordinator_task_name="traj_arm", - ) - traj_gen = MagicMock() - _install_generated_plan(module, config, traj_gen, [0.0, 1.0], [0.5, 1.5]) - module._state = ManipulationState.COMPLETED - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["arm/j1", "arm/j0"], position=[1.0, 0.0] - ) - mock_client = MagicMock() - module._coordinator_client = mock_client - - assert module.execute_plan() is False - - assert module._last_plan is not None - assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_not_called() - - def test_execute_plan_dispatches_selected_subsets_with_shared_clock_and_mapping(self): - left = RobotModelConfig( - name="left", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j0", "j1"], - planning_groups=[ - PlanningGroupDefinition( - name="wrist", joint_names=("j1",), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"left_coord_j1": "j1"}, - coordinator_task_name="traj_left", - ) - right = RobotModelConfig( - name="right", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["k0", "k1"], - planning_groups=[ - PlanningGroupDefinition( - name="elbow", joint_names=("k0",), base_link="base", tip_link="ee" - ) - ], - coordinator_task_name="traj_right", - ) - module = _make_module() - module._robots = { - "left": ("left_id", left, MagicMock()), - "right": ("right_id", right, MagicMock()), - } - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) - module._last_plan = GeneratedPlan( - group_ids=("left/wrist", "right/elbow"), - status=PlanningStatus.SUCCESS, - path=[ - JointState(name=["left/j1", "right/k0"], position=[0.0, 1.0]), - JointState(name=["left/j1", "right/k0"], position=[0.5, 1.5]), - ], - trajectory=JointTrajectory( - joint_names=["left/j1", "right/k0"], - points=[ - TrajectoryPoint( - time_from_start=0.0, - positions=[0.0, 1.0], - velocities=[0.0, 0.0], - ), - TrajectoryPoint( - time_from_start=2.5, - positions=[0.5, 1.5], - velocities=[0.2, 0.4], - ), - ], - ), - ) - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - module._world_monitor.get_current_joint_state.side_effect = lambda robot_id: { - "left_id": JointState(name=["j0", "j1"], position=[9.0, 0.0]), - "right_id": JointState(name=["k0", "k1"], position=[1.0, 9.0]), - }[robot_id] - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["left/j1", "right/k0"], position=[0.0, 1.0] - ) - - assert module.execute_plan() is True - - calls = mock_client.task_invoke.call_args_list - left_payload = calls[0].args[2]["trajectory"] - right_payload = calls[1].args[2]["trajectory"] - assert calls[0].args[0] == "traj_left" - assert left_payload.joint_names == ["left_coord_j1"] - assert [point.time_from_start for point in left_payload.points] == [0.0, 2.5] - assert [point.positions for point in left_payload.points] == [[0.0], [0.5]] - assert [point.velocities for point in left_payload.points] == [[0.0], [0.2]] - assert calls[1].args[0] == "traj_right" - assert right_payload.joint_names == ["k0"] - assert [point.time_from_start for point in right_payload.points] == [0.0, 2.5] - assert [point.positions for point in right_payload.points] == [[1.0], [1.5]] - assert [point.velocities for point in right_payload.points] == [[0.0], [0.4]] - def test_pose_wrappers_fail_safely_without_unique_pose_group(self, robot_config): no_pose_config = RobotModelConfig( name="test_arm", @@ -988,145 +431,10 @@ def test_pose_wrappers_fail_safely_with_multiple_pose_groups(self, robot_config) class TestJointNameTranslation: """Test trajectory joint name translation for coordinator.""" - def test_no_mapping_returns_original(self, robot_config, simple_trajectory): - """Without mapping, trajectory is returned unchanged.""" - module = _make_module() - - result = module._translate_trajectory_to_coordinator(simple_trajectory, robot_config) - assert result is simple_trajectory # Same object - - def test_mapping_translates_names(self, robot_config_with_mapping, simple_trajectory): - """With mapping, joint names are translated.""" - module = _make_module() - - result = module._translate_trajectory_to_coordinator( - simple_trajectory, robot_config_with_mapping - ) - assert result.joint_names == ["left/joint1", "left/joint2", "left/joint3"] - assert len(result.points) == 2 # Points preserved - class TestExecute: """Test coordinator execution.""" - def test_execute_requires_trajectory(self, robot_config): - """Execute fails without planned trajectory.""" - module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - - assert module.execute() is False - assert module._state == ManipulationState.IDLE - - def test_execute_requires_task_name(self): - """Execute fails without coordinator_task_name.""" - module = _make_module() - config_no_task = RobotModelConfig( - name="arm", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1"], - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1",), base_link="base_link", tip_link="ee" - ) - ], - ) - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory((0.0, [0.0]), (1.0, [0.1])) - _install_generated_plan(module, config_no_task, traj_gen, [0.0], [0.1]) - - assert module.execute("arm") is False - assert module._state == ManipulationState.IDLE - - def test_execute_plan_requires_task_name_without_sticking_executing(self, robot_config): - module = _make_module() - config_no_task = RobotModelConfig( - name="test_arm", - model_path=robot_config.model_path, - base_pose=robot_config.base_pose, - joint_names=robot_config.joint_names, - base_link=robot_config.base_link, - planning_groups=robot_config.planning_groups, - gripper_hardware_id="test_gripper", - ) - traj_gen = MagicMock() - traj_gen.generate.return_value = _make_trajectory( - (0.0, [0.0, 0.0, 0.0]), (1.0, [0.1, 0.2, 0.3]) - ) - module._robots = {"test_arm": ("id", config_no_task, traj_gen)} - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([config_no_task]) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._coordinator_client = MagicMock() - module._last_plan = GeneratedPlan( - trajectory=_generated_plan_trajectory( - ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - [0.0, 0.0, 0.0], - [0.1, 0.2, 0.3], - ), - group_ids=("test_arm/manipulator",), - status=PlanningStatus.SUCCESS, - path=[ - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ), - JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], - ), - ], - ) - - assert module.execute_plan() is False - assert module._state == ManipulationState.IDLE - - def test_execute_success(self, robot_config, simple_trajectory): - """Successful execute calls coordinator via task_invoke.""" - module = _make_module() - traj_gen = MagicMock() - traj_gen.generate.return_value = simple_trajectory - _install_generated_plan(module, robot_config, traj_gen, [0.0, 0.0, 0.0], [0.1, 0.2, 0.3]) - - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - - assert module.execute("test_arm") is True - assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_called_once() - task_name, method_name, payload = mock_client.task_invoke.call_args.args - assert task_name == "traj_arm" - assert method_name == "execute" - trajectory = payload["trajectory"] - assert trajectory.joint_names == simple_trajectory.joint_names - assert [point.positions for point in trajectory.points] == [ - [0.0, 0.0, 0.0], - [0.1, 0.2, 0.3], - ] - traj_gen.generate.assert_not_called() - - def test_execute_rejected(self, robot_config, simple_trajectory): - """Rejected execution sets FAULT state.""" - module = _make_module() - traj_gen = MagicMock() - traj_gen.generate.return_value = simple_trajectory - _install_generated_plan(module, robot_config, traj_gen, [0.0, 0.0, 0.0], [0.1, 0.2, 0.3]) - - mock_client = MagicMock() - mock_client.task_invoke.return_value = False - module._coordinator_client = mock_client - - assert module.execute("test_arm") is False - assert module._state == ManipulationState.FAULT - class TestRobotModelConfigMapping: """Test RobotModelConfig joint name mapping helpers.""" diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py index 935804834d..a056351c5f 100644 --- a/dimos/manipulation/test_plan_execution_reservation.py +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -1,33 +1,31 @@ # 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. -"""Tests for planned-path execution reservation and result propagation.""" +"""Module-side tests for runtime-owned plan submission and lifecycle handling.""" -from collections.abc import Callable +from collections.abc import Iterator +from inspect import signature from pathlib import Path -import threading -import time from unittest.mock import MagicMock -from dimos.manipulation._test_manipulation_helpers import make_module -from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationState +import pytest + +from dimos.manipulation._test_manipulation_helpers import ( + FakeCoordinatorGateway, + close_test_runtimes, + install_runtime, + make_module, +) +from dimos.manipulation.execution_runtime import LifecycleState, Outcome, prepare_generated_plan +from dimos.manipulation.manipulation_module import ( + ManipulationExecutionSnapshot, + ManipulationModule, +) from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus -from dimos.manipulation.planning.spec.models import GeneratedPlan, IKResult, PlanningResult -from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -36,26 +34,8 @@ from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -def _plan(group_id: str = "arm/manipulator") -> GeneratedPlan: - return GeneratedPlan( - group_ids=(group_id,), - path=[ - JointState(name=["arm/j0"], position=[0.0]), - JointState(name=["arm/j0"], position=[1.0]), - ], - trajectory=JointTrajectory( - joint_names=["arm/j0"], - points=[ - TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), - TrajectoryPoint(time_from_start=1.0, positions=[1.0], velocities=[0.0]), - ], - ), - ) - - -def _module_with_current(current: JointState) -> ManipulationModule: - module = make_module() - robot_config = RobotModelConfig( +def _config(task_name: str | None = "traj_arm") -> RobotModelConfig: + return RobotModelConfig( name="arm", model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), @@ -63,254 +43,148 @@ def _module_with_current(current: JointState) -> ManipulationModule: base_link="base", planning_groups=[ PlanningGroupDefinition( - name="manipulator", - joint_names=("j0",), - base_link="base", - tip_link="tool", + name="manipulator", joint_names=("j0",), base_link="base", tip_link="tool" ) ], max_velocity=1.0, max_acceleration=1.0, - coordinator_task_name="traj_arm", + coordinator_task_name=task_name, ) - module._robots = {"arm": ("robot_id", robot_config, MagicMock())} - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = current - module._world_monitor.is_state_stale.return_value = False - module._coordinator_client = MagicMock() - module._coordinator_client.task_invoke.return_value = True - return module -def _install_current_global_state(module: ManipulationModule, current: JointState) -> None: - robot_config = RobotModelConfig( - name="arm", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j0"], - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j0",), base_link="base", tip_link="tool" - ) - ], - coordinator_task_name="traj_arm", - ) - module._robots = {"arm": ("robot_id", robot_config, MagicMock())} - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = current - module._coordinator_client = MagicMock() - module._coordinator_client.task_invoke.return_value = True - - -def test_plan_to_pose_targets_propagates_failed_plan_as_false(): - module = make_module() - robot_config = RobotModelConfig( - name="test_arm", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["joint1", "joint2", "joint3"], - base_link="link_base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", - joint_names=("joint1", "joint2", "joint3"), - base_link="link_base", - tip_link="link_tcp", - ) +def _plan(group_id: str = "arm/manipulator") -> GeneratedPlan: + return GeneratedPlan( + group_ids=(group_id,), + path=[ + JointState(name=["arm/j0"], position=[0.0]), + JointState(name=["arm/j0"], position=[1.0]), ], - max_velocity=1.0, - max_acceleration=2.0, - coordinator_task_name="traj_arm", - ) - registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} - module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._kinematics = MagicMock() - module._kinematics.solve_pose_targets.return_value = IKResult( - status=IKStatus.SUCCESS, - joint_state=JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.1, 0.2, 0.3], + trajectory=JointTrajectory( + joint_names=["arm/j0"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), + TrajectoryPoint(time_from_start=1.0, positions=[1.0], velocities=[0.0]), + ], ), ) - module._planner = MagicMock() - module._planner.plan_selected_joint_path.return_value = PlanningResult( - status=PlanningStatus.NO_SOLUTION - ) - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - - assert module.plan_to_pose_targets({"test_arm/manipulator": pose}) is False - assert module._state == ManipulationState.FAULT -def test_execute_plan_selects_stored_plan_under_gate_and_reuses_it(): +def _module( + task_name: str | None = "traj_arm", +) -> tuple[ManipulationModule, FakeCoordinatorGateway]: module = make_module() + config = _config(task_name) + module._robots = {"arm": ("robot_id", config, object())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + gateway = install_runtime(module, [config]) + return module, gateway + + +@pytest.fixture(autouse=True) +def _close_runtime() -> Iterator[None]: + yield + close_test_runtimes() + + +def test_execution_snapshot_is_atomic_and_reports_ready_identity() -> None: + module, _gateway = _module() + runtime = module._execution_runtime + original_snapshot = runtime.snapshot + runtime.snapshot = MagicMock(wraps=original_snapshot) + + view = module.get_execution_snapshot() + assert isinstance(view, ManipulationExecutionSnapshot) + assert view.state == LifecycleState.IDLE.name + assert view.diagnostic == "" + assert view.ready_plan_id is None + assert not view.has_ready_plan + runtime.snapshot.assert_called_once_with() + + runtime.snapshot.reset_mock() + token = runtime.start_planning() + prepared = prepare_generated_plan(_plan(), module._execution_topology) + assert runtime.complete_planning(token, prepared).accepted + runtime.snapshot.reset_mock() + view = module.get_execution_snapshot() + assert view.state == LifecycleState.READY.name + assert view.ready_plan_id is not None + assert view.has_ready_plan + runtime.snapshot.assert_called_once_with() + + +def test_execution_snapshot_is_safe_before_runtime_initialization() -> None: + view = make_module().get_execution_snapshot() + assert view == ManipulationExecutionSnapshot("IDLE", "", None, False) + + +def test_execution_api_has_no_robot_name_and_uses_exact_ready_identity() -> None: + module, gateway = _module() plan = _plan() - module._last_plan = plan - module._state = ManipulationState.COMPLETED - _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) - - assert module.execute_plan() is True - assert module.execute_plan() is True - assert module._coordinator_client.task_invoke.call_count == 2 + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(plan, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + assert "robot_name" not in signature(module.execute).parameters + assert "robot_name" not in signature(module.execute_plan).parameters + assert module.execute_plan(plan=plan) is True + assert len(gateway.execute_calls) == 1 + assert module.get_state() in (LifecycleState.RUNNING.name, LifecycleState.IDLE.name) + assert module.cancel() is True + + +def test_trajectory_status_uses_operation_handle_identity() -> None: + module, _gateway = _module() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(_plan(), module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + result = module._execution_runtime.execute_ready() + assert result.accepted and result.value is not None + operation = module._execution_runtime.snapshot().operation + assert operation is not None + status = module.get_trajectory_status() + assert status is not None + assert status["operation_id"] == operation.handle.operation_id + assert module.cancel() is True + + +def test_execute_plan_explicit_plan_prepares_without_using_other_ready_plan() -> None: + module, gateway = _module() + ready = _plan() + explicit = _plan() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(ready, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + assert module.execute_plan(plan=explicit) is True + assert len(gateway.execute_calls) == 1 + assert module.cancel() is True + + +def test_plan_without_coordinator_task_fails_during_preparation_without_dispatch() -> None: + module, gateway = _module(task_name=None) + assert module.execute_plan(plan=_plan()) is False + assert gateway.execute_calls == [] + assert module.get_state() == LifecycleState.IDLE.name -def test_execute_and_execute_plan_race_without_consuming_latest_plan(): - module = make_module() +def test_cancel_ready_clears_plan_and_status_is_snapshot_only() -> None: + module, _gateway = _module() plan = _plan() - module._last_plan = plan - module._state = ManipulationState.COMPLETED - _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) - callers_ready = threading.Barrier(3) - dispatch_started = threading.Event() - allow_dispatch_return = threading.Event() - results: list[bool] = [] - - def dispatch(*_args, **_kwargs): - dispatch_started.set() - assert allow_dispatch_return.wait(timeout=1.0) - return True - - def call_execute(execute: Callable[[], bool]) -> None: - callers_ready.wait(timeout=1.0) - results.append(execute()) - - module._coordinator_client.task_invoke.side_effect = dispatch - legacy = threading.Thread(target=call_execute, args=(module.execute,)) - current = threading.Thread(target=call_execute, args=(module.execute_plan,)) - legacy.start() - current.start() - callers_ready.wait(timeout=1.0) - assert dispatch_started.wait(timeout=1.0) - time.sleep(0.05) - - allow_dispatch_return.set() - legacy.join(timeout=1.0) - current.join(timeout=1.0) - - assert not legacy.is_alive() - assert not current.is_alive() - assert sorted(results) == [False, True] - assert module._coordinator_client.task_invoke.call_count == 1 - - -def test_execute_started_during_cancel_cannot_dispatch_after_cancel(): - module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) - module._last_plan = _plan() - module._state = ManipulationState.COMPLETED - module._possibly_active_tasks = {"traj_arm"} - cancel_started = threading.Event() - execute_finished = threading.Event() - release_cancel = threading.Event() - cancel_result: list[bool] = [] - execute_result: list[bool] = [] - - def invoke(task_name: str, method: str, _payload: dict[str, object]) -> bool: - assert task_name == "traj_arm" - if method == "cancel": - cancel_started.set() - assert release_cancel.wait(timeout=1.0) - return False - assert method == "execute" - raise AssertionError("execution dispatched during cancellation") - - module._coordinator_client.task_invoke.side_effect = invoke - cancelling = threading.Thread(target=lambda: cancel_result.append(module.cancel())) - cancelling.start() - assert cancel_started.wait(timeout=1.0) - - def run_execute() -> None: - execute_result.append(module.execute_plan()) - execute_finished.set() - - executing = threading.Thread(target=run_execute) - executing.start() - assert execute_finished.wait(timeout=1.0) - assert execute_result == [False] - assert not any( - call.args[1] == "execute" for call in module._coordinator_client.task_invoke.call_args_list - ) - release_cancel.set() - cancelling.join(timeout=1.0) - executing.join(timeout=1.0) - - assert cancel_result == [True] - module._coordinator_client.task_invoke.assert_called_once_with("traj_arm", "cancel", {}) - - -def test_execute_plan_accepts_a_direct_plan_without_reserving_stored_plan(): - module = make_module() - direct_plan = _plan() - module._last_plan = _plan("other/manipulator") - _install_current_global_state(module, JointState(name=["arm/j0"], position=[0.0])) + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(plan, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + assert module.get_trajectory_status() is not None + assert module.cancel() is True + assert module._execution_runtime.snapshot().ready_plan is None - assert module.execute_plan(plan=direct_plan) is True - module._coordinator_client.task_invoke.assert_called_once() - - -def test_execute_plan_uses_current_global_state_without_consuming_cached_plan(): - module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) - cached_plan = _plan("cached/manipulator") - module._last_plan = cached_plan - assert module.execute_plan(plan=_plan()) is True - module._coordinator_client.task_invoke.assert_called_once() - assert module._last_plan is cached_plan - - -def test_execute_plan_rejects_missing_direct_plan_start_before_dispatch_without_reserving_cached_plan(): - module = _module_with_current(JointState(name=[], position=[])) - cached_plan = _plan("cached/manipulator") - module._last_plan = cached_plan +def test_reset_fault_and_runtime_shutdown_are_owned_by_runtime() -> None: + module, gateway = _module() + gateway.execute_outcome = Outcome.REJECTED assert module.execute_plan(plan=_plan()) is False - module._coordinator_client.task_invoke.assert_not_called() - assert module._last_plan is cached_plan - - -def test_execute_plan_rejects_malformed_direct_plan_before_dispatch_without_reserving_cached_plan(): - module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) - cached_plan = _plan("cached/manipulator") - direct_plan = _plan() - direct_plan.trajectory.points[0].positions = [] - module._last_plan = cached_plan - - assert module.execute_plan(plan=direct_plan) is False - module._coordinator_client.task_invoke.assert_not_called() - assert module._last_plan is cached_plan - - -def test_execute_plan_pre_dispatch_exception_restores_previous_state_without_faulting(): - module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) - module._last_plan = _plan() - module._state = ManipulationState.COMPLETED - module._prepare_execution = MagicMock(side_effect=RuntimeError("split exploded")) - - assert module.execute_plan() is False - - assert module._state == ManipulationState.COMPLETED - assert module._error_message == "Failed to prepare execution: split exploded" - module._coordinator_client.task_invoke.assert_not_called() - - -def test_execute_plan_dispatch_exception_faults_without_sticking_executing(): - module = _module_with_current(JointState(name=["arm/j0"], position=[0.0])) - module._last_plan = _plan() - module._coordinator_client.task_invoke.side_effect = RuntimeError("rpc exploded") - - assert module.execute_plan() is False - - assert module._state == ManipulationState.FAULT - assert module._error_message.startswith("Failed to dispatch trajectory: rpc exploded") - assert "traj_arm" in module._error_message + assert module.get_state() == LifecycleState.IDLE.name + assert module.reset().is_success() + assert module.get_state() == LifecycleState.IDLE.name + module.stop() diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py index 21af8dfd1b..77f707c084 100644 --- a/dimos/manipulation/visualization/operator.py +++ b/dimos/manipulation/visualization/operator.py @@ -19,7 +19,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field import math -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID, RobotName @@ -33,11 +33,21 @@ @dataclass(frozen=True) class OperatorStatus: - """Compact dynamic manipulation status.""" + """One coherent projection of the manipulation runtime snapshot.""" state: str - error: str - has_plan: bool + diagnostic: str = "" + ready_plan_status: str = "NONE" + ready_plan_id: str | None = None + # These two names are retained for callers that predate the snapshot API. + error: str | None = None + has_plan: bool | None = None + + def __post_init__(self) -> None: + if self.error is None: + object.__setattr__(self, "error", self.diagnostic) + if self.has_plan is None: + object.__setattr__(self, "has_plan", self.ready_plan_status == "READY") @dataclass(frozen=True) @@ -79,11 +89,14 @@ def __init__(self, module: ManipulationModule, world_monitor: WorldMonitor) -> N self._world_monitor = world_monitor def status(self) -> OperatorStatus: - """Return compact dynamic state without topology or joint telemetry.""" + """Return a coherent runtime snapshot without topology or telemetry.""" + snapshot: Any = self._module.get_execution_snapshot() + state_name = getattr(snapshot.state, "name", str(snapshot.state)) return OperatorStatus( - state=self._module.get_state(), - error=self._module.get_error(), - has_plan=self._module.has_planned_path(), + state=state_name, + diagnostic=snapshot.diagnostic or "", + ready_plan_status="READY" if snapshot.has_ready_plan else "NONE", + ready_plan_id=(None if snapshot.ready_plan_id is None else str(snapshot.ready_plan_id)), ) def get_init_joints(self, robot_name: RobotName) -> JointState | None: diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py index 701e6625da..a924a173c2 100644 --- a/dimos/manipulation/visualization/test_operator.py +++ b/dimos/manipulation/visualization/test_operator.py @@ -15,6 +15,9 @@ """Focused tests for the manipulation visualization operator facade.""" from pathlib import Path +from types import SimpleNamespace + +import pytest from dimos.agents.skill_result import SkillResult from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition @@ -76,7 +79,7 @@ def _robot_config( class FakeModule: def __init__(self) -> None: - self.state = "COMPLETED" + self.state = "IDLE" self.error = "" self.has_plan = True self.plan = GeneratedPlan( @@ -111,6 +114,14 @@ def __init__(self) -> None: def get_state(self) -> str: return self.state + def get_execution_snapshot(self) -> SimpleNamespace: + return SimpleNamespace( + state=self.state, + diagnostic=self.error, + has_ready_plan=self.has_plan, + ready_plan_id="fake-plan" if self.has_plan else None, + ) + def get_error(self) -> str: return self.error @@ -251,7 +262,7 @@ def test_status_is_compact_and_does_not_read_topology_or_telemetry() -> None: status = operator.status() - assert status.state == "COMPLETED" + assert status.state == "IDLE" assert status.error == "" assert status.has_plan is True assert module.topology_calls == 0 @@ -259,6 +270,25 @@ def test_status_is_compact_and_does_not_read_topology_or_telemetry() -> None: assert monitor.telemetry_calls == 0 +@pytest.mark.parametrize( + "state", + ("IDLE", "PLANNING", "READY", "DISPATCHING", "RUNNING", "CANCELLING", "FAULT"), +) +def test_status_projects_every_runtime_lifecycle_state_from_one_snapshot(state: str) -> None: + operator, module, _ = _operator() + module.state = state + module.has_plan = state == "READY" + module.error = "runtime fault" if state == "FAULT" else "" + + status = operator.status() + + assert status.state == state + assert status.diagnostic == ("runtime fault" if state == "FAULT" else "") + assert status.has_plan is (state == "READY") + assert status.ready_plan_status == ("READY" if state == "READY" else "NONE") + assert status.ready_plan_id == ("fake-plan" if state == "READY" else None) + + def test_evaluate_joint_target_accepts_exact_global_selection_domain() -> None: operator, _, _ = _operator() diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index e959f82f76..a71536f2b9 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -23,6 +23,7 @@ from dimos.manipulation.visualization.operator import ( JointTargetRequest, ManipulationOperator, + OperatorStatus, PoseTargetRequest, TargetEvaluationResult, ) @@ -232,10 +233,16 @@ def get_module_state(self) -> str: return self.operator.status().state def get_error(self) -> str: - return self.operator.status().error + return self.operator.status().error or "" def reset(self) -> bool: - return self.operator.reset() + result = self.operator.reset() + if result: + self.state.error = "" + self.state.action_status = ActionStatus.IDLE + self.state.runtime_failure = False + self.state.local_action_failure = False + return result def evaluate_joint_target_set( self, group_ids: Sequence[PlanningGroupID], targets: Mapping[PlanningGroupID, JointState] @@ -291,6 +298,7 @@ def plan_to_selected_joints( JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) ) self.state.plan_state.plan = plan + self.state.plan_state.plan_id = None return plan is not None def preview_path(self) -> bool: @@ -304,6 +312,7 @@ def execute(self) -> bool: def refresh(self) -> None: if self._closed: return + status = self.operator.status() robots = self.list_robots() groups = self.list_planning_groups() self.state.backend_status = ( @@ -320,7 +329,8 @@ def refresh(self) -> None: if set(self.state.group_joint_targets) != initialized_groups: self._build_joint_sliders() self._sync_group_selector(groups) - self._refresh_selected_robot_state() + self._refresh_selected_robot_state(status) + self._apply_operator_status(status) self._ensure_scene_controls() self._sync_target_ghost_visibility() self._sync_preset_dropdown() @@ -440,18 +450,56 @@ def _set_scene_grid_visible(self, visible: bool) -> None: return self.scene.set_reference_grid_visible(bool(visible)) - def _refresh_selected_robot_state(self) -> None: + def _refresh_selected_robot_state(self, status: OperatorStatus) -> None: robot_name = self.state.selected_robot if robot_name is None: self.state.current_joints = None - self.state.manipulation_state = self.get_module_state() + self.state.manipulation_state = status.state return current = self.get_current_joint_state(robot_name) self.state.current_joints = list(current.position) if current is not None else None - self.state.manipulation_state = self.get_module_state() - adapter_error = self.get_error() - if adapter_error: - self.state.error = adapter_error + self.state.manipulation_state = status.state + + def _apply_operator_status(self, status: OperatorStatus) -> None: + """Project one operator snapshot onto the panel state.""" + self.state.manipulation_state = status.state + self.state.ready_plan_id = getattr(status, "ready_plan_id", None) + snapshot_status = hasattr(status, "diagnostic") + diagnostic = getattr(status, "diagnostic", getattr(status, "error", "")) + if diagnostic: + self.state.error = diagnostic + elif snapshot_status and self.state.action_status != ActionStatus.FAILED: + self.state.error = "" + if status.state == "FAULT": + self.state.runtime_failure = True + self.state.action_status = ActionStatus.FAILED + elif status.state == "CANCELLING": + self.state.action_status = ActionStatus.CANCELLING + elif status.state in {"DISPATCHING", "RUNNING"}: + self.state.action_status = ActionStatus.EXECUTING + elif status.state == "PLANNING" and self.state.action_status == ActionStatus.IDLE: + self.state.action_status = ActionStatus.RUNNING + elif status.state in {"IDLE", "READY"}: + was_runtime_failure = self.state.runtime_failure + self.state.runtime_failure = False + if self.state.action_status in { + ActionStatus.RUNNING, + ActionStatus.EXECUTING, + ActionStatus.CANCELLING, + }: + self.state.action_status = ActionStatus.IDLE + elif was_runtime_failure and not self.state.local_action_failure: + self.state.action_status = ActionStatus.IDLE + if not diagnostic: + self.state.error = "" + if status.state != "READY" and self.state.plan_state.status == PlanStatus.FRESH: + if status.state in {"DISPATCHING", "RUNNING", "CANCELLING", "IDLE", "FAULT"}: + self.state.plan_state.status = PlanStatus.STALE + if ( + self.state.plan_state.status == PlanStatus.FRESH + and self.state.plan_state.plan_id != self.state.ready_plan_id + ): + self.state.plan_state.status = PlanStatus.STALE def _ensure_scene_controls(self) -> None: if self.scene is None: @@ -1093,9 +1141,11 @@ def operation() -> None: ) return if ok: + status = self.operator.status() self.state.plan_state.status = PlanStatus.FRESH self.state.plan_state.group_ids = group_ids self.state.plan_state.target_sequence_id = target_sequence_id + self.state.plan_state.plan_id = getattr(status, "ready_plan_id", None) else: self.state.plan_state.status = PlanStatus.FAILED self.state.plan_state.plan = None @@ -1172,7 +1222,6 @@ def operation() -> None: ) return self.state.action_status = ActionStatus.EXECUTING - self.state.plan_state.status = PlanStatus.EXECUTING ok = self.execute() if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): self._finish_operation( @@ -1181,9 +1230,13 @@ def operation() -> None: return if not ok: self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation( - f"execute={ok}", operation_id=operation_id, selection_epoch=selection_epoch - ) + self._finish_operation( + f"execute={ok}", operation_id=operation_id, selection_epoch=selection_epoch + ) + else: + # True means dispatch was accepted, not that motion completed. + self.state.last_result = "execute=True (dispatched)" + self.refresh() self._operation_worker.submit( operation, on_error=lambda message: self._set_operation_error(message, operation_id) @@ -1209,10 +1262,7 @@ def _submit_cancel(self) -> None: def _mark_cancelled_plan_state(self, cancelled_action: ActionStatus) -> None: if self.state.plan_state.status == PlanStatus.PLANNING: self.state.plan_state.status = PlanStatus.FAILED - elif ( - cancelled_action == ActionStatus.EXECUTING - or self.state.plan_state.status == PlanStatus.EXECUTING - ): + elif cancelled_action == ActionStatus.EXECUTING: self.state.plan_state.status = PlanStatus.STALE def _restart_operation_worker(self) -> None: @@ -1272,6 +1322,7 @@ def _finish_operation( self.state.action_status = ActionStatus.IDLE if clear_error: self.state.error = "" + self.state.local_action_failure = False self.state.last_result = result self.refresh() @@ -1283,13 +1334,16 @@ def _set_operation_error(self, message: str, operation_id: int) -> None: def _set_recoverable_error(self, message: str) -> None: if self._closed: return - self.state.error = message self.refresh() + # This is a panel-local validation message, not a runtime diagnostic. + self.state.error = message + self._update_status_text() def _set_error(self, message: str) -> None: if self._closed: return self.state.action_status = ActionStatus.FAILED + self.state.local_action_failure = True self.state.error = message self.refresh() diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index 95a21e5b84..4e8090f229 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -66,7 +66,6 @@ class PlanStatus(str, Enum): PLANNING = "planning" FRESH = "fresh" STALE = "stale" - EXECUTING = "executing" FAILED = "failed" @@ -97,6 +96,7 @@ class PanelPlanState: group_ids: tuple[PlanningGroupID, ...] = () target_sequence_id: int = 0 plan: GeneratedPlan | None = None + plan_id: str | None = None @dataclass @@ -112,7 +112,10 @@ class PanelState: backend_status: BackendConnectionStatus = BackendConnectionStatus.DISCONNECTED target_status: TargetStatus = TargetStatus.EMPTY action_status: ActionStatus = ActionStatus.IDLE + runtime_failure: bool = False + local_action_failure: bool = False manipulation_state: str = "DISCONNECTED" + ready_plan_id: str | None = None current_joints: list[float] | None = None cartesian_target: Pose | None = None feasibility: FeasibilityState = field(default_factory=FeasibilityState) @@ -146,7 +149,7 @@ def can_plan(self) -> bool: and bool(self.selected_group_ids) and self.action_status == ActionStatus.IDLE and self.target_status == TargetStatus.FEASIBLE - and self.manipulation_state in {"IDLE", "COMPLETED"} + and self.manipulation_state in {"IDLE", "READY"} and self.plan_state.status != PlanStatus.PLANNING ) @@ -155,15 +158,24 @@ def can_preview(self) -> bool: self.runtime == PanelRuntime.RUNNING and self.backend_status == BackendConnectionStatus.READY and self.action_status == ActionStatus.IDLE + and self.manipulation_state == "READY" and self.plan_state.status == PlanStatus.FRESH + and self.plan_state.plan_id is not None + and self.ready_plan_id is not None + and self.plan_state.plan_id == self.ready_plan_id ) def can_cancel(self) -> bool: - return self.action_status in { + return self.manipulation_state in { + "PLANNING", + "DISPATCHING", + "RUNNING", + "CANCELLING", + } or self.action_status in { ActionStatus.RUNNING, ActionStatus.PREVIEWING, ActionStatus.EXECUTING, - } or (self.manipulation_state == "EXECUTING") + } def can_execute(self, action_status: ActionStatus | None = None) -> bool: plan = self.plan_state @@ -173,9 +185,12 @@ def can_execute(self, action_status: ActionStatus | None = None) -> bool: and self.backend_status == BackendConnectionStatus.READY and effective_action_status == ActionStatus.IDLE and self.target_status == TargetStatus.FEASIBLE - and self.manipulation_state in {"IDLE", "COMPLETED"} + and self.manipulation_state == "READY" and plan.status == PlanStatus.FRESH and plan.plan is not None + and plan.plan_id is not None + and self.ready_plan_id is not None + and plan.plan_id == self.ready_plan_id and plan.group_ids == self.selected_group_ids and plan.target_sequence_id == self.latest_sequence_id ): diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index 9d8be1e168..0116821a81 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -58,7 +58,9 @@ def __init__(self, module: FakeOperatorBackend | None = None) -> None: self.module = module or FakeOperatorBackend() def status(self) -> OperatorStatus: - return OperatorStatus(state="IDLE", error="", has_plan=False) + return OperatorStatus( + state="READY", diagnostic="", ready_plan_status="READY", ready_plan_id="fake-plan" + ) def get_init_joints(self, robot_name: str) -> None: return None @@ -273,7 +275,10 @@ def test_gui_only_preview_submits_timeout_override(monkeypatch: pytest.MonkeyPat monkeypatch.setattr(gui, "_operation_worker", FakeTimeoutSubmitWorker(submissions)) gui.state.runtime = PanelRuntime.RUNNING gui.state.backend_status = BackendConnectionStatus.READY + gui.state.manipulation_state = "READY" + gui.state.ready_plan_id = "fake-plan" gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.plan_id = "fake-plan" gui._submit_preview() assert submissions == [{"timeout_seconds": 0.25}] @@ -290,11 +295,13 @@ def test_gui_preview_enters_previewing_before_worker_runs( gui.state.runtime = PanelRuntime.RUNNING gui.state.backend_status = BackendConnectionStatus.READY gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "COMPLETED" + gui.state.manipulation_state = "READY" gui.state.selected_group_ids = ("arm/manipulator",) gui.state.plan_state.status = PlanStatus.FRESH gui.state.plan_state.group_ids = gui.state.selected_group_ids gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, trajectory=JointTrajectory(), @@ -327,11 +334,13 @@ def test_gui_selection_change_clears_invalidated_preview( gui.state.runtime = PanelRuntime.RUNNING gui.state.backend_status = BackendConnectionStatus.READY gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "COMPLETED" + gui.state.manipulation_state = "READY" gui.state.selected_group_ids = (groups[0].id,) gui.state.plan_state.status = PlanStatus.FRESH gui.state.plan_state.group_ids = gui.state.selected_group_ids gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, trajectory=JointTrajectory(), @@ -363,11 +372,13 @@ def test_gui_selection_change_ignores_invalidated_preview_error( gui.state.runtime = PanelRuntime.RUNNING gui.state.backend_status = BackendConnectionStatus.READY gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "COMPLETED" + gui.state.manipulation_state = "READY" gui.state.selected_group_ids = (groups[0].id,) gui.state.plan_state.status = PlanStatus.FRESH gui.state.plan_state.group_ids = gui.state.selected_group_ids gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" gui.state.plan_state.plan = GeneratedPlan( group_ids=gui.state.selected_group_ids, trajectory=JointTrajectory(), @@ -470,3 +481,46 @@ def test_gui_ignores_stale_timed_out_operation_finish() -> None: assert gui.state.action_status == ActionStatus.FAILED assert gui.state.error == "Operation timed out after 5.0s" + + +def test_gui_projects_all_runtime_states_and_fault_reset_preserves_local_failure() -> None: + gui = make_gui() + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.selected_group_ids = ("arm/manipulator",) + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.action_status = ActionStatus.IDLE + + expected_actions = { + "IDLE": ActionStatus.IDLE, + "PLANNING": ActionStatus.RUNNING, + "READY": ActionStatus.IDLE, + "DISPATCHING": ActionStatus.EXECUTING, + "RUNNING": ActionStatus.EXECUTING, + "CANCELLING": ActionStatus.CANCELLING, + "FAULT": ActionStatus.FAILED, + } + for state, expected_action in expected_actions.items(): + gui._apply_operator_status( + OperatorStatus( + state=state, + diagnostic="fault" if state == "FAULT" else "", + ready_plan_status="READY" if state == "READY" else "NONE", + ready_plan_id="fake-plan" if state == "READY" else None, + ) + ) + assert gui.state.manipulation_state == state + assert gui.state.action_status == expected_action + + gui._apply_operator_status(OperatorStatus("IDLE")) + assert gui.state.can_plan() is True + + gui.state.action_status = ActionStatus.FAILED + gui.state.local_action_failure = True + gui._apply_operator_status(OperatorStatus("FAULT", diagnostic="fault")) + gui._apply_operator_status( + OperatorStatus("READY", ready_plan_status="READY", ready_plan_id="fake-plan") + ) + + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.local_action_failure is True diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index c4c638374f..52d201d9f4 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -29,7 +29,7 @@ from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningSceneInfo -from dimos.manipulation.visualization.operator import TargetEvaluationResult +from dimos.manipulation.visualization.operator import OperatorStatus, TargetEvaluationResult from dimos.manipulation.visualization.viser.animation import ( GroupPreviewAnimation, PreviewFrame, @@ -203,6 +203,7 @@ def __init__(self, groups: list[PlanningGroup], states: dict[str, JointState]) - self.cancelled = 0 self.cleared = 0 self.last_plan: GeneratedPlan | None = None + self.state = "IDLE" def make_plan(self, group_ids: tuple[str, ...]) -> GeneratedPlan: names = [ @@ -223,6 +224,7 @@ def make_plan(self, group_ids: tuple[str, ...]) -> GeneratedPlan: status=PlanningStatus.SUCCESS, ) self.last_plan = plan + self.state = "READY" return plan def list_robots(self) -> list[str]: @@ -244,7 +246,7 @@ def get_init_joints(self, name: str) -> JointState: return JointState({"name": self.configs[name].joint_names, "position": [-0.5, -1.0]}) def get_state(self) -> str: - return "IDLE" + return self.state def get_error(self) -> str: return "" @@ -261,6 +263,7 @@ def preview_plan(self) -> bool: def execute(self) -> bool: self.executions += 1 + self.state = "DISPATCHING" return True def cancel(self) -> bool: @@ -300,11 +303,12 @@ def __init__(self, module: Module, monitor: Monitor) -> None: self.module = module self.monitor = monitor - def status(self) -> SimpleNamespace: - return SimpleNamespace( + def status(self) -> OperatorStatus: + return OperatorStatus( state=self.module.get_state(), - error=self.module.get_error(), - has_plan=True, + diagnostic=self.module.get_error(), + ready_plan_status="READY", + ready_plan_id="fake-plan", ) def get_init_joints(self, robot_name: str) -> JointState | None: @@ -532,6 +536,7 @@ def test_plan_target_sequence_invalidation_and_unfiltered_all_robot_execute( group_ids=(left.id, right.id), target_sequence_id=gui.state.latest_sequence_id, plan=module.last_plan, + plan_id="fake-plan", ) gui.state.target_status = TargetStatus.FEASIBLE gui._submit_execute() diff --git a/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml b/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml new file mode 100644 index 0000000000..29e56d8d3b --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-20 diff --git a/openspec/changes/extract-manipulation-execution-runtime/design.md b/openspec/changes/extract-manipulation-execution-runtime/design.md new file mode 100644 index 0000000000..46d26cd201 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/design.md @@ -0,0 +1,115 @@ +# Design: Manipulation Execution Runtime + +## Context + +`ManipulationModule` currently owns planning, coordinator dispatch, cancellation, remote-task monitoring, and lifecycle bookkeeping. These responsibilities are interleaved with manually assigned state and separate execution flags, so concurrent requests can produce ambiguous dispatch, completion, or cancellation outcomes. Remote coordinator activity is also difficult to reconcile safely when a response is delayed, incomplete, or contradictory. + +## Baseline inventory (historical, task 1.1) + +The pre-runtime implementation used `ManipulationState.IDLE`, `PLANNING`, `EXECUTING`, `COMPLETED`, and `FAULT`, with module-owned `_state`, `_last_plan`, `_execution_generation`, `_planning_epoch`, `_execution_dispatch_lock`, and `_coordinator_client` fields. Planning entered through `plan_to_pose`, `plan_to_joints`, `plan_to_pose_targets`, and `plan_to_joint_targets`; execution entered through `execute`/`execute_plan` and split the stored plan into per-robot trajectories. + +Module-owned coordinator RPC sites included trajectory `execute`, `cancel`, and `get_state` calls, plus fallback completion polling and direct coordinator-client setup. Status consumers read `get_state`, `get_error`, and `get_trajectory_status`, while the operator/UI adapters consumed those module methods. + +The current replacement seams are `ExecutionRuntime` and its `LifecycleState`, `RuntimeSnapshot`, plan/operation/task records, and serialized `start_planning`/`complete_planning`/`execute_ready`/`execute_explicit`/`cancel`/`reset`/`poll` methods. The four planning APIs now submit planning epochs/results through those planning methods, and execution entry points submit prepared or ready plans through the runtime. `ControlCoordinatorGateway` now owns the runtime-facing `execute`/`cancel`/`status`/`reset` RPC boundary. `ManipulationModule.get_state`, `get_execution_snapshot`, `get_error`, `cancel`, `reset`, and execution entry points delegate to that snapshot/runtime seam; operator/UI status reads the atomic snapshot rather than legacy module flags. + +This change introduces an internal execution runtime owned by `ManipulationModule`. The runtime is the sole lifecycle owner for manipulation execution. The module remains responsible for running planning algorithms and submitting planning and execution events to the runtime; it does not independently assign lifecycle state or maintain a second execution lifecycle. + +The runtime provides the authoritative lifecycle and status snapshot for the manipulation layer while encapsulating coordinator trajectory-task interaction. The design is internal and implementation-agnostic: it specifies ownership, events, state semantics, and safety rules rather than a particular threading or queue library. + +## Goals / Non-Goals + +### Goals + +- Define one authoritative lifecycle with exactly these public states: `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT`. +- Serialize planning results, execute, cancel, reset, remote reconciliation, and status publication through a single-owner queue and thread (or equivalent single serialized owner). +- Reject execution requests received during an active attempt rather than queueing them for later dispatch. +- Use a formal reducer to derive lifecycle state and runtime data from events, rather than ad hoc state assignments. +- Make the runtime authoritative for the ready plan, operation IDs, attempts, task knowledge, faults, and the coordinator client. +- Make plans single-use and make replacement of a ready plan explicit and atomic. +- Dispatch every robot trajectory represented by a generated plan; caller-selected partial-robot execution is not supported. +- Reconcile coordinator task activity conservatively, latching `FAULT` whenever safe remote activity cannot be proven. +- Require reset to establish that relevant remote activity is inactive before leaving `FAULT`. +- Build and validate coordinator-name inverse mapping once, using canonical configuration references. + +### Non-Goals + +- This change does not redesign planning algorithms or trajectory representations. +- It does not add a replacement for manipulation-layer trajectory start-state freshness validation. +- It does not introduce partial-robot execution or retain a duplicate direct trajectory-execution path owned by `ManipulationModule`. Standalone diagnostic/operator CLI direct-RPC tooling in `dimos/manipulation/control/coordinator_client.py` is out of scope; moving or removing it requires a separate compatibility-scoped change. +- It does not define a new external runtime or coordinator dependency. +- It does not prescribe a UI layout, transport, or concrete queue/thread implementation. + +## Decisions + +### Runtime ownership and event processing + +The internal runtime is the sole lifecycle owner. `ManipulationModule` invokes planning algorithms and submits events such as planning started, plan produced, execute requested, cancel requested, reset requested, and coordinator observations. The module must not mutate lifecycle state, ready-plan ownership, operation bookkeeping, or fault state outside the runtime. + +All lifecycle-affecting work is serialized by a single-owner queue and thread, or an equivalent mechanism with the same single-writer guarantee. This includes execute, cancel, reset, coordinator polling, task reconciliation, and status snapshot publication. External callers may enqueue requests concurrently, but event order at the runtime owner is the order used for lifecycle decisions. + +Execution admission is state-based: while the runtime is `DISPATCHING`, `RUNNING`, or `CANCELLING`, a new execution request is completed as terminal `REJECTED` and is not retained for later dispatch. A coordinator refusal before a remote attempt is established is also a terminal `REJECTED` result and returns the runtime to `IDLE`; neither case is a safety fault. + +A formal reducer is the authority for applying events. It validates each event against the current state and runtime context, performs the associated side effects through the runtime-owned coordinator client where needed, and produces the next state plus an atomic status snapshot. Invalid or unsafe events do not create an alternate state path; they are rejected or become a fault according to the event's safety semantics. + +### Public lifecycle + +The public state vocabulary is limited to: + +| State | Meaning | +| --- | --- | +| `IDLE` | No ready plan or active manipulation operation is owned by the runtime. | +| `PLANNING` | A planning request is in progress; its eventual result is not yet executable. | +| `READY` | One complete plan is owned by the runtime and may be executed once. | +| `DISPATCHING` | The selected plan has been consumed and coordinator submission/reconciliation is in progress. | +| `RUNNING` | The current operation is known to be active remotely. | +| `CANCELLING` | Cancellation has been requested and the runtime is reconciling until inactivity or fault is established. | +| `FAULT` | The runtime cannot prove a safe, unambiguous lifecycle or remote-task condition; execution remains latched until reset succeeds. | + +The snapshot is atomic from consumers' perspective and includes the public state together with the applicable operation, attempt, plan, task, and fault information. Consumers must use the snapshot rather than infer lifecycle from independent flags. + +### Runtime-owned data and coordinator access + +The runtime owns the ready plan, operation IDs, attempt records, task knowledge, fault details, and coordinator client. An operation ID identifies a logical execution request; attempts identify coordinator submission/reconciliation efforts within that operation. This separates local request identity from remote task identity and prevents stale responses from being treated as current activity. + +The coordinator client used for `ManipulationModule` execution is accessed only by the runtime. A plan is dispatched through the runtime's single coordinator path; active module-owned duplicate direct trajectory execution paths are not retained. This does not require moving or removing `dimos/manipulation/control/coordinator_client.py`, which remains standalone diagnostic/operator CLI direct-RPC tooling; changing it is a separate compatibility-scoped change. The runtime reconciles submission responses and subsequent task observations against the current operation and attempt before changing state. + +### Plan replacement and single-use semantics + +Plans are single-use. A plan selected for dispatch is consumed before dispatch begins, so it cannot be executed again or reused after a failed or completed attempt. Planning while a ready plan exists may produce an explicit replacement event; applying that event atomically discards the prior ready plan and installs the new complete plan. There is never an externally observable interval in which both plans are executable. + +Only a plan containing every robot trajectory represented by its planning groups is accepted. Partial-robot selection and a `robot_name` execution-selection argument are not part of this lifecycle. A plan that omits a trajectory required by its own planning groups is rejected rather than converted into a partial execution. + +### Coordinator-name mapping + +The coordinator-name inverse mapping is constructed once during runtime initialization from canonical configuration references. Initialization validates that references are complete, names are unambiguous, and the mapping is internally consistent. Runtime operation uses this validated mapping and does not rebuild or infer it from mutable per-request data. + +### Freshness and uncertainty safety + +Manipulation-layer trajectory start-state freshness validation is removed in this change and has no replacement here. The runtime instead focuses on lifecycle and remote-task certainty. + +Remote uncertainty is fail-safe and latching. If coordinator task activity, ownership, submission outcome, cancellation outcome, or completion cannot be proven safe and attributable to the current operation, the reducer enters `FAULT`. A reset request may clear the fault only after the runtime proves that relevant remote activity is inactive; reset is not merely a local flag clear. Until then, the runtime remains in `FAULT` and does not dispatch a new plan. + +## Risks / Trade-offs + +- A single serialized owner simplifies race handling and makes reducer behavior deterministic, but polling, coordinator calls, or planning-result delivery can delay subsequent lifecycle events if not bounded or isolated appropriately. +- Conservative latching can expose `FAULT` for recoverable coordinator/network ambiguity and requires operator-visible reset, trading availability for avoidance of duplicate or uncontrolled motion. +- Consuming a plan before dispatch prevents accidental reuse but means a failed dispatch cannot silently retry the same plan; any retry must be represented as a new, explicit attempt under the runtime's rules. +- Atomic replacement makes plan ownership clear, but a newly completed plan can supersede a prior ready plan without execution of the prior one. +- Building the inverse mapping once improves consistency and detects configuration errors early, but configuration changes require runtime reinitialization rather than being picked up dynamically. +- Removing freshness validation leaves start-state correctness to lower layers or future work; this change must not imply that freshness is guaranteed elsewhere. + +## Migration Plan + +1. Add the internal runtime behind `ManipulationModule` and route all lifecycle-affecting requests, planning results, coordinator polling, cancellation, reset, and status publication through its event interface. +2. Move ready-plan storage, operation and attempt bookkeeping, task knowledge, fault handling, coordinator access, and state transitions into the runtime. Remove parallel module-level lifecycle flags and assignments. +3. Update manipulation state adapters and callers to consume the seven public states and atomic snapshots. Remove the `robot_name` execution-selection argument and any partial-robot execution paths. +4. Construct and validate the coordinator-name inverse mapping once from canonical configuration references during runtime initialization. +5. Remove manipulation-layer trajectory start-state freshness validation and active `ManipulationModule` duplicate direct trajectory execution paths, without adding replacement validation in this change. Leave `dimos/manipulation/control/coordinator_client.py` unchanged; moving or removing that standalone diagnostic/operator CLI direct-RPC tooling is a separate compatibility-scoped change. +6. Validate normal planning, atomic replacement, single-use dispatch, completion, cancellation, stale/ambiguous coordinator observations, and reset-after-inactivity behavior before removing the legacy lifecycle implementation. + +## Open Questions + +- What concrete polling interval and coordinator-call timeout best balance responsiveness with the runtime's single-owner serialization? +- Which coordinator observations constitute sufficient proof of inactivity for each reset path, and how long should that proof remain valid? +- How should faults and operation/attempt identifiers be exposed in existing UI and operator-facing diagnostics beyond the atomic snapshot? +- What future layer, if any, will own trajectory start-state freshness validation after it is removed from manipulation? diff --git a/openspec/changes/extract-manipulation-execution-runtime/proposal.md b/openspec/changes/extract-manipulation-execution-runtime/proposal.md new file mode 100644 index 0000000000..0eeedc65be --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/proposal.md @@ -0,0 +1,29 @@ +## Why + +`ManipulationModule` currently combines planning, coordinator dispatch, cancellation, task-status monitoring, and lifecycle bookkeeping. Its manually assigned state, separate execution flags, and ambiguous dispatch/completion semantics make concurrent cancellation and remote-task uncertainty difficult to reason about safely. + +## What Changes + +- Extract coordinator trajectory dispatch and its full remote-task lifecycle into an internal execution runtime owned by `ManipulationModule`. +- Replace manual lifecycle assignments with a single-owner, event-driven state machine and atomic status snapshots. +- Replace the public lifecycle vocabulary with `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT`. +- Make plans single-use: replacement planning discards the prior ready plan, and dispatch consumes the selected plan. +- Serialize execute, cancel, reset, and status-polling work through the runtime; latch faults when coordinator task activity cannot be proven safe. +- Remove manipulation-layer trajectory start-state freshness validation without adding a replacement in this change. +- Remove partial-robot execution and the duplicate direct trajectory execution path. + +## Capabilities + +### New Capabilities +- `manipulation-lifecycle-runtime`: Defines the authoritative manipulation plan and execution lifecycle, state transitions, snapshots, and fault/reset behavior. +- `coordinator-trajectory-execution`: Defines safe plan-to-coordinator dispatch, task reconciliation, completion monitoring, and cancellation semantics. + +### Modified Capabilities + +None. + +## Impact + +- Affects `dimos/manipulation/manipulation_module.py`, manipulation operator/UI state adapters, and manipulation execution tests. +- Changes public state strings and removes the `robot_name` execution-selection argument; callers and UI state handling must be updated. +- Centralizes use of `ControlCoordinator` trajectory task RPCs in the internal runtime; no new external dependency is introduced. diff --git a/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md b/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md new file mode 100644 index 0000000000..d1c7ea1f5a --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md @@ -0,0 +1,85 @@ +## ADDED Requirements + +### Requirement: Serialized coordinator task execution +The coordinator trajectory execution runtime SHALL serialize execute, cancel, reset, and status-polling operations. It SHALL allow at most one active plan and one active execution attempt at a time, and SHALL NOT issue planning or execution work while a remote execution is active. + +The runtime SHALL reject a second execution request while an execution attempt is active; execution requests SHALL NOT be queued for later dispatch. The rejection SHALL be terminal for the rejected request and SHALL not fault or interrupt the already active attempt. + +#### Scenario: Active execution rejects a second request +- **WHEN** an execution request arrives while the first attempt is active +- **THEN** the second request SHALL complete with a terminal `REJECTED` outcome without being queued +- **AND** the active attempt SHALL remain unchanged +- **AND** the coordinator SHALL receive no execute request for the rejected operation + +#### Scenario: Rejected execution does not fault the idle runtime +- **WHEN** a valid execution request reaches dispatch and the coordinator rejects it +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** the manipulation lifecycle SHALL return to `IDLE` +- **AND** the runtime SHALL NOT transition to `FAULT` + +### Requirement: Global-to-local-to-coordinator mapping boundary +The runtime SHALL define an explicit mapping boundary from the global manipulation plan to the local coordinator trajectory/task request and back to the global manipulation result. The mapping SHALL preserve plan identity, attempt identity, trajectory/action data, and terminal outcome without exposing coordinator-specific task representation as the manipulation lifecycle contract. + +The runtime SHALL perform coordinator dispatch only from a fully mapped local request and SHALL map every coordinator result into the normalized global outcome before publishing it to manipulation callers. + +#### Scenario: Global plan is mapped for coordinator dispatch +- **WHEN** a global plan is selected for execution +- **THEN** the runtime SHALL create local coordinator requests containing every trajectory represented by the global plan and their correlation identities +- **AND** it SHALL dispatch that local request rather than passing an unbounded global plan object through the coordinator API + +#### Scenario: Coordinator completion is mapped back +- **WHEN** the coordinator reports a terminal result for the active task +- **THEN** the runtime SHALL map it to the corresponding global plan/attempt outcome +- **AND** callers SHALL receive the normalized manipulation outcome rather than raw coordinator task details + +### Requirement: Register the task before execute RPC +The runtime MUST register the active plan, attempt, and coordinator task correlation context before invoking the coordinator execute RPC. It MUST associate the RPC response with that pre-registered attempt and MUST retain the context when dispatch is accepted. + +#### Scenario: Accepted task is correlated without a race +- **WHEN** the runtime is about to invoke the coordinator execute RPC +- **THEN** its active-attempt/task registration SHALL already exist +- **AND** an immediate acceptance or completion response SHALL be attributed to that attempt + +### Requirement: Normalize cancellation and unknown outcomes +The runtime SHALL normalize coordinator cancellation responses into a cancellation outcome. It SHALL normalize unknown, malformed, unavailable, or contradictory coordinator task states into a safety-uncertain outcome rather than treating them as successful completion or ordinary cancellation. + +Cancellation SHALL be idempotent for an already terminal attempt, but cancellation of an active task SHALL not be reported complete until the coordinator confirms cancellation or confirms that the task is no longer active. If that confirmation cannot be obtained, the runtime SHALL latch remote safety uncertainty and SHALL block new execution until reset reconciliation. + +#### Scenario: Active cancellation waits for confirmation +- **WHEN** cancellation is requested for an active coordinator task +- **THEN** the runtime SHALL serialize the cancel request and enter cancellation handling +- **AND** it SHALL report cancellation only after the task is confirmed cancelled or inactive + +#### Scenario: Unknown task state is unsafe +- **WHEN** coordinator polling returns an unknown task state or cannot establish whether the task remains active +- **THEN** the runtime SHALL normalize the result as safety uncertainty +- **AND** it SHALL latch the runtime in `FAULT` until safe reset reconciliation + +### Requirement: Reset reconciles remote safety +Reset SHALL first reconcile coordinator task activity and SHALL clear local plan, attempt, and pending request state only after remote safety is proven. Reset SHALL leave the runtime faulted when remote activity is still active or cannot be determined, and SHALL be retryable without issuing a new trajectory execution. + +#### Scenario: Reset clears local state after safe reconciliation +- **WHEN** reset confirms that the coordinator has no active task +- **THEN** the runtime SHALL clear the active attempt, ready plan, and pending execution-request bookkeeping +- **AND** it SHALL return the manipulation lifecycle to `IDLE` + +#### Scenario: Reset does not hide active remote execution +- **WHEN** reset finds an active coordinator task +- **THEN** the runtime SHALL not clear the safety fault or start another execution +- **AND** it SHALL continue reconciliation or require cancellation before becoming idle + +### Requirement: Plan-complete execution only +The coordinator execution interface SHALL require every robot trajectory represented by the generated plan and SHALL not accept a partial-robot argument or robot-name selector. The runtime SHALL maintain a single trajectory execution path through the coordinator and SHALL remove duplicate direct execution paths owned by `ManipulationModule`. This requirement does not remove or relocate `dimos/manipulation/control/coordinator_client.py`, whose standalone diagnostic/operator CLI direct-RPC tooling is outside the `ManipulationModule` execution path; changing that tooling is a separate compatibility-scoped change. + +#### Scenario: Partial execution input has no dispatch effect +- **WHEN** an execution request omits a robot trajectory represented by its generated plan or supplies a robot selector +- **THEN** the runtime SHALL reject the request as invalid +- **AND** it SHALL issue no coordinator execute RPC + +### Requirement: No manipulation start-state freshness gate +The coordinator execution runtime SHALL not add or restore manipulation-layer validation that rejects a plan because its trajectory start state is not fresh. Start-state acceptance and physical safety remain the responsibility of the coordinator and its execution path. + +#### Scenario: Execution is not rejected by the removed freshness gate +- **WHEN** a complete mapped trajectory reaches the coordinator boundary with a non-fresh manipulation start-state observation +- **THEN** the runtime SHALL still submit the coordinator request +- **AND** it SHALL resolve the result according to coordinator acceptance, task status, cancellation, or safety uncertainty diff --git a/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md b/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md new file mode 100644 index 0000000000..89f1ff09b7 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md @@ -0,0 +1,96 @@ +## ADDED Requirements + +### Requirement: Authoritative lifecycle state +The manipulation execution runtime SHALL own the authoritative lifecycle state and SHALL expose only `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT` as lifecycle states. + +The runtime SHALL publish an atomic status snapshot containing the lifecycle state and the currently relevant plan, attempt, task, and fault information. Readers SHALL observe one internally consistent snapshot rather than independently read lifecycle flags. + +#### Scenario: Snapshot reports a coherent active execution +- **WHEN** a plan has been accepted for execution and its coordinator task is active +- **THEN** one status snapshot SHALL report `RUNNING` together with the associated plan/attempt and remote task identity +- **AND** the snapshot SHALL NOT combine fields from different plans or attempts + +#### Scenario: Invalid lifecycle transition is rejected +- **WHEN** an event requests a transition that is not legal from the current lifecycle state +- **THEN** the runtime SHALL leave the lifecycle state unchanged +- **AND** it SHALL report the rejected transition without silently changing an execution flag + +### Requirement: Legal lifecycle transitions +The runtime SHALL permit planning only from `IDLE` or `READY`, transitioning to `PLANNING`. A successful planning result SHALL transition to `READY`; a planning failure SHALL return to `IDLE` unless the runtime is already faulted. Dispatch of a selected ready plan SHALL transition through `DISPATCHING` and, after accepted dispatch, to `RUNNING`. A completed or cancelled attempt SHALL transition to `IDLE`. Cancellation of an active attempt SHALL transition to `CANCELLING` before reaching `IDLE` or `FAULT`. A safety-uncertain condition SHALL transition to and latch `FAULT`. + +The runtime SHALL NOT transition out of `FAULT` except through reset reconciliation that proves the remote task state is safe. + +#### Scenario: Replacement planning discards a ready plan +- **WHEN** the runtime is `READY` and a new planning request is accepted +- **THEN** the lifecycle SHALL become `PLANNING` +- **AND** the previous ready plan SHALL be discarded +- **AND** only the newly produced plan MAY become `READY` + +#### Scenario: Completed execution returns to idle +- **WHEN** the active attempt is physically reported complete +- **THEN** the runtime SHALL consume the attempt and transition to `IDLE` +- **AND** its next snapshot SHALL contain no executable ready plan + +### Requirement: Planning and single-use ready plans +The runtime SHALL serialize planning with dispatch, cancellation, reset, and status polling. It SHALL maintain at most one ready plan. Planning SHALL produce a plan with a unique plan identity and SHALL make it executable only after the planning operation succeeds. Dispatch SHALL consume the selected plan exactly once; a plan SHALL NOT be dispatched again after dispatch acceptance, completion, cancellation, or failure. + +The runtime SHALL reject planning and dispatch requests while an attempt is actively executing remotely (`DISPATCHING`, `RUNNING`, or `CANCELLING`), except that cancellation and reset operations SHALL remain available through their defined serialized paths. + +An execution request rejected because an attempt is active SHALL complete with a terminal `REJECTED` outcome, SHALL NOT be queued, and SHALL NOT transition the active runtime to `FAULT`. + +#### Scenario: Dispatch consumes a plan +- **WHEN** a ready plan is dispatched and the coordinator accepts the task +- **THEN** the plan SHALL no longer be available as a ready plan +- **AND** a repeated dispatch request using that plan identity SHALL be rejected + +#### Scenario: Planning is blocked during remote execution +- **WHEN** the runtime is `RUNNING` or `CANCELLING` +- **THEN** a planning request SHALL be rejected without creating or replacing a plan +- **AND** no trajectory execution request SHALL be issued + +#### Scenario: Execution is rejected during remote execution +- **WHEN** the runtime is `DISPATCHING`, `RUNNING`, or `CANCELLING` and a second execution request is received +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** it SHALL remain unqueued and issue no coordinator execute request +- **AND** the active lifecycle state SHALL remain unchanged + +#### Scenario: Coordinator rejection returns to idle +- **WHEN** a valid execution request is rejected by the coordinator before an active remote attempt is established +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** the lifecycle SHALL return to `IDLE` +- **AND** the runtime SHALL NOT enter `FAULT` + +### Requirement: Dispatch acceptance is distinct from physical completion +The runtime SHALL distinguish coordinator dispatch acceptance from physical trajectory completion. Dispatch acceptance SHALL establish an active attempt and `RUNNING` state, but SHALL NOT be reported as physical completion. The runtime SHALL transition to completion only after coordinator task status or an equivalent authoritative result reports that the robot physically completed the task. + +#### Scenario: Accepted dispatch remains running +- **WHEN** the coordinator accepts a trajectory task but has not reported physical completion +- **THEN** the runtime SHALL report `RUNNING` +- **AND** it SHALL continue task monitoring +- **AND** it SHALL NOT report the manipulation as complete + +### Requirement: Remote safety uncertainty is latched +If the runtime cannot prove that a coordinator task is safely inactive or reconciled, it SHALL latch a `FAULT` state and SHALL prevent new planning or execution. The fault SHALL remain latched across subsequent polls and SHALL be cleared only by reset reconciliation that establishes a safe remote state. + +#### Scenario: Unknown remote activity faults the runtime +- **WHEN** a status, cancellation, or reset operation returns an unknown/error result that does not prove the remote task is inactive +- **THEN** the runtime SHALL transition to `FAULT` +- **AND** subsequent plan and execute requests SHALL be rejected + +#### Scenario: Safe reset clears a latched fault +- **WHEN** reset reconciliation confirms that no coordinator task is active and the remote state is safe +- **THEN** the runtime SHALL clear the latched fault +- **AND** it SHALL transition to `IDLE` with no ready plan or active attempt + +### Requirement: No manipulation-layer start-state or partial-robot execution contract +The manipulation runtime SHALL NOT validate trajectory start-state freshness as a prerequisite for manipulation execution. It SHALL NOT accept a partial-robot execution selection or a `robot_name` execution argument; manipulation execution SHALL target every robot trajectory represented by the generated plan through the coordinator boundary. + +#### Scenario: Stale manipulation start state does not block dispatch +- **WHEN** a ready manipulation plan is otherwise valid but its trajectory start state is not fresh according to the removed manipulation-layer check +- **THEN** the runtime SHALL not reject the plan solely for that freshness condition +- **AND** execution SHALL proceed subject to coordinator acceptance and safety rules + +#### Scenario: Partial robot selection is rejected +- **WHEN** a caller attempts manipulation execution with a partial-robot selection or robot-specific execution argument +- **THEN** the request SHALL be rejected +- **AND** no coordinator trajectory request SHALL be issued diff --git a/openspec/changes/extract-manipulation-execution-runtime/tasks.md b/openspec/changes/extract-manipulation-execution-runtime/tasks.md new file mode 100644 index 0000000000..179c30cf00 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/tasks.md @@ -0,0 +1,71 @@ +## 1. Characterize existing behavior and update contracts + +- [x] 1.1 Inventory current manipulation lifecycle states, execution flags, coordinator RPC call sites, planning APIs, and status consumers; see the historical baseline below. +- [x] 1.2 Add characterization tests for current planning, dispatch acceptance, physical completion, cancellation, reset, and coordinator uncertainty behavior. +- [x] 1.3 Define the runtime event, normalized outcome, operation/attempt identity, and atomic snapshot APIs in the manipulation execution test contract. +- [x] 1.4 Update API tests and call sites to use `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT` only. +- [x] 1.5 Update API tests to reject `robot_name` and all partial-robot execution selections. + +## 2. Implement lifecycle state, reducer, and snapshots + +- [x] 2.1 Add the lifecycle state model with exactly the seven public states and explicit legal-transition validation. +- [x] 2.2 Add runtime event types for planning, plan replacement, execute, cancel, reset, coordinator responses, task observations, and polling. +- [x] 2.3 Implement the formal reducer so invalid transitions leave state unchanged and produce a rejected-transition result. +- [x] 2.4 Implement atomic status snapshots containing consistent state, plan, operation, attempt, task, and fault data. +- [x] 2.5 Add reducer tests for legal transitions, rejected transitions, coherent snapshots, and dispatch acceptance versus physical completion. + +## 3. Introduce the single-owner execution runtime + +- [x] 3.1 Add the internal runtime owned by `ManipulationModule` with a single-owner event queue and serialized processing loop. +- [x] 3.2 Move coordinator client ownership, ready-plan storage, operation IDs, attempts, task knowledge, and fault details into the runtime. +- [x] 3.3 Ensure execute, cancel, reset, planning results, coordinator polling, and snapshot publication all use the serialized runtime path. +- [x] 3.4 Add runtime tests proving that concurrent requests cannot create overlapping execution attempts or inconsistent snapshots. +- [x] 3.5 Add tests proving that planning and execution work are blocked while a remote attempt is active, while cancellation and reset remain available. + +## 4. Move planning results and ready-plan lifecycle + +- [x] 4.1 Route planning start requests from `ManipulationModule` to the runtime and transition accepted requests to `PLANNING`. +- [x] 4.2 Route planning success and failure results to runtime events with unique plan identities and normalized outcomes. +- [x] 4.3 Implement atomic replacement so a new successful planning request discards the prior ready plan before exposing the new plan as `READY`. +- [x] 4.4 Enforce validation that every robot trajectory represented by a plan is present, and reject incomplete plans without creating executable ready state. +- [x] 4.5 Implement single-use plan consumption before dispatch and reject repeated dispatch of a consumed plan identity. +- [x] 4.6 Add tests for planning from `IDLE` and `READY`, replacement, planning failure, plan identity, and single-use dispatch semantics. + +## 5. Implement mapping, dispatch, and task knowledge + +- [x] 5.1 Build and validate the coordinator-name inverse mapping once from canonical configuration references during runtime initialization. +- [x] 5.2 Add the global-plan to local-coordinator-request mapping and normalized coordinator-result mapping while preserving plan and attempt identities. +- [x] 5.3 Register active plan, operation, attempt, and task-correlation context before invoking the coordinator execute RPC. +- [x] 5.4 Implement the single complete-robot coordinator dispatch path through `DISPATCHING` to `RUNNING` on acceptance. +- [x] 5.5 Record coordinator task knowledge and reject stale or mismatched responses without attributing them to the active attempt. +- [x] 5.6 Add tests for mapping validation, complete trajectory requests, pre-registration, immediate responses, accepted-but-not-complete tasks, and normalized terminal results. + +## 6. Implement cancellation, polling, fault, and reset safety + +- [x] 6.1 Implement serialized status polling and task reconciliation for active attempts, including physical completion detection. +- [x] 6.2 Implement cancellation through `CANCELLING`, reporting cancellation only after coordinator cancellation or inactivity is confirmed. +- [x] 6.3 Normalize unknown, malformed, unavailable, and contradictory coordinator results as safety uncertainty and latch `FAULT`. +- [x] 6.4 Ensure latched `FAULT` rejects new planning and execution and remains latched across subsequent polls. +- [x] 6.5 Implement reset reconciliation that clears plan, attempt, queue, and fault state only after remote inactivity is proven, returning to `IDLE`. +- [x] 6.6 Ensure reset remains retryable and does not issue new trajectory execution while remote activity is active or uncertain. +- [x] 6.7 Add tests for rejection of execution during an active attempt (without queueing), terminal coordinator rejection returning to `IDLE`, active cancellation confirmation, unknown task states, fault latching, and safe/unsafe reset. + +## 7. Wire module, operator, and UI compatibility + +- [x] 7.1 Replace direct lifecycle assignments and execution flags in `ManipulationModule` with runtime event submission and snapshot consumption. +- [x] 7.2 Update manipulation operators and callers for the new execution API, normalized outcomes, and removal of robot-specific selection. +- [x] 7.3 Update UI/state adapters to consume atomic snapshots and map all seven public states without legacy state names. +- [x] 7.4 Add compatibility tests covering module status publication, operator actions, UI state mapping, fault visibility, and reset behavior. + +## 8. Remove legacy paths and validate the migration + +- [x] 8.1 Remove legacy module-level lifecycle bookkeeping, duplicate coordinator access, and active `ManipulationModule` direct trajectory execution paths. Do not move or remove `dimos/manipulation/control/coordinator_client.py`: it is standalone diagnostic/operator CLI direct-RPC tooling, not `ManipulationModule` execution; changing it is a separate compatibility-scoped change. +- [x] 8.2 Remove manipulation-layer trajectory start-state freshness validation without adding a replacement gate in this change. +- [x] 8.3 Remove partial-robot execution handling and the `robot_name` argument from implementation, adapters, and tests. +- [x] 8.4 Run targeted lifecycle runtime, coordinator execution, manipulation module, operator, and UI compatibility tests. +- [x] 8.5 Run formatting, lint, type checks, and the relevant OpenSpec or repository validation checks. Verified with `uv run --with mypy mypy dimos/manipulation`: 54 source files, no issues found. +- [x] 8.6 Review test coverage for every lifecycle state, transition, safety fault, reset path, plan replacement, and single-use invariant. + +## Verification note + +Self-hosted and Drake-dependent verification is environment-limited for this change and is not claimed unless run on the corresponding self-hosted/Drake environment. Local artifact review and environment-independent checks may be run separately; unavailable self-hosted/Drake checks must be reported as not run rather than inferred as passing. From d74ced30443d7a78c4ef72e054792d5b5b007032 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 22 Jul 2026 11:57:38 -0700 Subject: [PATCH 29/29] refactor(manipulation): simplify execution runtime --- .../_test_manipulation_helpers.py | 38 +- dimos/manipulation/execution_auxiliary.py | 70 + dimos/manipulation/execution_clock.py | 32 + dimos/manipulation/execution_effects.py | 113 + dimos/manipulation/execution_gateway.py | 105 + dimos/manipulation/execution_models.py | 188 + dimos/manipulation/execution_policy.py | 134 + dimos/manipulation/execution_runtime.py | 3156 +++++++---------- dimos/manipulation/execution_topology.py | 273 ++ dimos/manipulation/manipulation_module.py | 336 +- .../manipulation/test_execution_auxiliary.py | 13 + dimos/manipulation/test_execution_clock.py | 82 + dimos/manipulation/test_execution_effects.py | 91 + dimos/manipulation/test_execution_gateway.py | 59 + dimos/manipulation/test_execution_models.py | 27 + dimos/manipulation/test_execution_policy.py | 116 + dimos/manipulation/test_execution_runtime.py | 2803 ++++++++------- dimos/manipulation/test_execution_topology.py | 199 ++ .../manipulation/test_manipulation_module.py | 158 +- .../test_plan_execution_reservation.py | 122 +- dimos/manipulation/visualization/viser/gui.py | 6 + .../manipulation/visualization/viser/state.py | 4 - .../visualization/viser/test_gui.py | 36 + .../visualization/viser/test_state.py | 16 +- .../.openspec.yaml | 2 + .../design.md | 111 + .../proposal.md | 35 + .../coordinator-trajectory-execution/spec.md | 51 + .../manipulation-lifecycle-runtime/spec.md | 71 + .../tasks.md | 71 + 30 files changed, 5031 insertions(+), 3487 deletions(-) create mode 100644 dimos/manipulation/execution_auxiliary.py create mode 100644 dimos/manipulation/execution_clock.py create mode 100644 dimos/manipulation/execution_effects.py create mode 100644 dimos/manipulation/execution_gateway.py create mode 100644 dimos/manipulation/execution_models.py create mode 100644 dimos/manipulation/execution_policy.py create mode 100644 dimos/manipulation/execution_topology.py create mode 100644 dimos/manipulation/test_execution_auxiliary.py create mode 100644 dimos/manipulation/test_execution_clock.py create mode 100644 dimos/manipulation/test_execution_effects.py create mode 100644 dimos/manipulation/test_execution_gateway.py create mode 100644 dimos/manipulation/test_execution_models.py create mode 100644 dimos/manipulation/test_execution_policy.py create mode 100644 dimos/manipulation/test_execution_topology.py create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/design.md create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/proposal.md create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md create mode 100644 openspec/changes/simplify-manipulation-execution-runtime/tasks.md diff --git a/dimos/manipulation/_test_manipulation_helpers.py b/dimos/manipulation/_test_manipulation_helpers.py index 35d089cb58..006fa3bf0b 100644 --- a/dimos/manipulation/_test_manipulation_helpers.py +++ b/dimos/manipulation/_test_manipulation_helpers.py @@ -16,9 +16,14 @@ from unittest.mock import MagicMock -from dimos.manipulation.execution_runtime import ExecutionRuntime, Outcome +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + Outcome, + prepare_generated_plan, +) from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan _TEST_RUNTIMES: list[ExecutionRuntime] = [] @@ -26,8 +31,17 @@ class FakeCoordinatorGateway: """Deterministic coordinator gateway used by module-side tests.""" - def __init__(self, execute_outcome: Outcome = Outcome.ACCEPTED) -> None: + def __init__( + self, + execute_outcome: Outcome = Outcome.ACCEPTED, + status_outcome: Outcome = Outcome.INACTIVE, + status_outcomes: list[Outcome] | None = None, + cancel_outcome: Outcome = Outcome.CANCELLED, + ) -> None: self.execute_outcome = execute_outcome + self.status_outcome = status_outcome + self.status_outcomes = status_outcomes or [] + self.cancel_outcome = cancel_outcome self.execute_calls: list[tuple[str, object]] = [] self.cancel_calls: list[str] = [] self.stopped = False @@ -38,11 +52,13 @@ def execute(self, task_name: str, request: object) -> Outcome: def cancel(self, task_name: str) -> Outcome: self.cancel_calls.append(task_name) - return Outcome.CANCELLED + return self.cancel_outcome def status(self, task_name: str) -> Outcome: del task_name - return Outcome.INACTIVE + if self.status_outcomes: + return self.status_outcomes.pop(0) + return self.status_outcome def reset(self, task_name: str) -> Outcome: del task_name @@ -71,7 +87,7 @@ def __init__(self) -> None: self._execution_runtime = None self._execution_topology = None self._init_joints = {} - self.config = MagicMock(planning_timeout=10.0) + self.config = MagicMock(planning_timeout=10.0, physical_operation_timeout=60.0) def make_module() -> ManipulationModule: @@ -93,12 +109,24 @@ def install_runtime( lambda: gateway, topology=module._execution_topology, action_timeout=module.config.planning_timeout, + physical_operation_timeout=module.config.physical_operation_timeout, poll_interval=0.01, ) _TEST_RUNTIMES.append(module._execution_runtime) return gateway +def install_ready_plan(module: ManipulationModule, plan: GeneratedPlan) -> None: + """Publish a prepared plan through the production runtime API.""" + assert module._execution_runtime is not None + assert module._execution_topology is not None + token = module._execution_runtime.start_planning() + assert isinstance(token, str) + prepared = prepare_generated_plan(plan, module._execution_topology) + result = module._execution_runtime.complete_planning(token, prepared) + assert result.accepted + + def close_test_runtimes() -> None: """Close runtimes created by the lightweight test harness.""" while _TEST_RUNTIMES: diff --git a/dimos/manipulation/execution_auxiliary.py b/dimos/manipulation/execution_auxiliary.py new file mode 100644 index 0000000000..650e0e5f71 --- /dev/null +++ b/dimos/manipulation/execution_auxiliary.py @@ -0,0 +1,70 @@ +"""Passive auxiliary-call bookkeeping for the execution runtime.""" + +from dataclasses import dataclass +from typing import Any + +from dimos.manipulation.execution_effects import AuxiliaryDone + + +@dataclass(frozen=True) +class AuxiliaryTicket: + action_id: str + setter: bool + + +class AuxiliaryCallBook: + """Store auxiliary correlations and results without lifecycle decisions.""" + + def __init__(self) -> None: + self._pending: dict[str, float] = {} + self._inflight: set[str] = set() + self._results: dict[str, tuple[bool, Any, str]] = {} + + def register(self, action_id: str, deadline: float, setter: bool) -> AuxiliaryTicket: + self._inflight.add(action_id) + self._pending[action_id] = deadline + return AuxiliaryTicket(action_id, setter) + + def complete(self, done: AuxiliaryDone) -> bool: + pending = done.action_id in self._pending + inflight = done.action_id in self._inflight + self._pending.pop(done.action_id, None) + self._inflight.discard(done.action_id) + if not pending and not inflight: + return False + if done.action_id in self._results: + return True + if done.error is None: + result = (True, done.value, "") + else: + result = (False, None, str(done.error)) + self._results[done.action_id] = result + if len(self._results) > 16: + self._results.pop(next(iter(self._results))) + return True + + def take_result(self, action_id: str) -> tuple[bool, Any, str] | None: + return self._results.pop(action_id, None) + + def has_pending(self) -> bool: + return bool(self._pending) + + def has_inflight(self) -> bool: + return bool(self._inflight) + + def deadlines(self) -> tuple[float, ...]: + return tuple(self._pending.values()) + + def has_unsettled(self) -> bool: + return bool(self._pending or self._inflight) + + def expire(self, now: float, diagnostic: str) -> bool: + changed = False + for action_id, deadline in tuple(self._pending.items()): + if deadline <= now and action_id not in self._results: + self._pending.pop(action_id, None) + self._results[action_id] = (False, None, diagnostic) + changed = True + if len(self._results) > 16: + self._results.pop(next(iter(self._results))) + return changed diff --git a/dimos/manipulation/execution_clock.py b/dimos/manipulation/execution_clock.py new file mode 100644 index 0000000000..1ac6d4a587 --- /dev/null +++ b/dimos/manipulation/execution_clock.py @@ -0,0 +1,32 @@ +"""Validated monotonic clock dependency for execution deadlines.""" + +from collections.abc import Callable +import math +import threading + + +class InvalidMonotonicClock(ValueError): # noqa: N818 + """Raised when a clock violates the execution clock contract.""" + + +class ValidatedMonotonicClock: + """Serialize sampling and validate finite, nondecreasing seconds.""" + + def __init__(self, source: Callable[[], float]) -> None: + self._source = source + self._lock = threading.Lock() + with self._lock: + initial = self._source() + if not math.isfinite(initial): + raise InvalidMonotonicClock("monotonic clock returned a non-finite value") + self._last = initial + + def now(self) -> float: + with self._lock: + sample = self._source() + if not math.isfinite(sample): + raise InvalidMonotonicClock("monotonic clock returned a non-finite value") + if sample < self._last: + raise InvalidMonotonicClock("monotonic clock moved backwards") + self._last = sample + return sample diff --git a/dimos/manipulation/execution_effects.py b/dimos/manipulation/execution_effects.py new file mode 100644 index 0000000000..18ccaf8d98 --- /dev/null +++ b/dimos/manipulation/execution_effects.py @@ -0,0 +1,113 @@ +"""Daemon worker plumbing for execution-runtime effects.""" + +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures.thread import _worker +from dataclasses import dataclass +import threading +from typing import Any, cast +import weakref + + +@dataclass(frozen=True) +class EffectDone: + action_id: str + outcome: Any + + +@dataclass(frozen=True) +class AuxiliaryDone: + action_id: str + value: Any + error: BaseException | None = None + + +@dataclass(frozen=True) +class StopDone: + success: bool + diagnostic: str = "" + + +class _DaemonThreadPoolExecutor(ThreadPoolExecutor): + def _adjust_thread_count(self) -> None: + if self._idle_semaphore.acquire(timeout=0): + return + if len(self._threads) >= self._max_workers: + return + + def weakref_cb(_ref: Any, q: Any = self._work_queue) -> None: + q.put(None) + + thread = threading.Thread( + name=f"{self._thread_name_prefix}_{len(self._threads)}", + target=_worker, + args=( + weakref.ref(self, weakref_cb), + self._work_queue, + self._initializer, # type: ignore[attr-defined] + self._initargs, # type: ignore[attr-defined] + ), + daemon=True, + ) + thread.start() + cast("Any", self._threads).add(thread) + + +class ExecutionEffectRunner: + """Execute already-admitted effects and post immutable completion events.""" + + def __init__(self) -> None: + self._executor = _DaemonThreadPoolExecutor( + max_workers=4, thread_name_prefix="execution-rpc" + ) + + def submit_action( + self, action_id: str, effect: Callable[[], Any], enqueue: Callable[[Any], None] + ) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + outcome = future.exception() + if outcome is None: + outcome = future.result() + except BaseException as exc: + outcome = exc + enqueue(EffectDone(action_id, outcome)) + + future.add_done_callback(done_callback) + + def submit_auxiliary( + self, action_id: str, effect: Callable[[], Any], enqueue: Callable[[Any], None] + ) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + error = future.exception() + enqueue( + AuxiliaryDone( + action_id, + None if error is not None else future.result(), + error, + ) + ) + except BaseException as exc: + enqueue(AuxiliaryDone(action_id, None, exc)) + + future.add_done_callback(done_callback) + + def submit_stop(self, effect: Callable[[], Any], enqueue: Callable[[Any], None]) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + error = future.exception() + enqueue(StopDone(error is None, "" if error is None else str(error))) + except BaseException as exc: + enqueue(StopDone(False, str(exc))) + + future.add_done_callback(done_callback) + + def shutdown(self) -> None: + self._executor.shutdown(wait=False, cancel_futures=False) diff --git a/dimos/manipulation/execution_gateway.py b/dimos/manipulation/execution_gateway.py new file mode 100644 index 0000000000..b439d26859 --- /dev/null +++ b/dimos/manipulation/execution_gateway.py @@ -0,0 +1,105 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Coordinator RPC conversion seam for the manipulation execution runtime.""" + +from __future__ import annotations + +from typing import Any + +from dimos.manipulation.execution_models import Outcome +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +class ControlCoordinatorGateway: + """Convert coordinator RPC requests and responses to normalized outcomes.""" + + def __init__(self, client: Any) -> None: + self._client = client + + def _invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + return self._client.task_invoke(task, method, kwargs) + + def execute(self, task_name: str, request: Any) -> Outcome: + try: + value = self._invoke( + task_name, + "execute", + request if isinstance(request, dict) else {"trajectory": request}, + ) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def cancel(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "cancel", {}) + return ( + Outcome.CANCELLED + if value is True + else Outcome.INACTIVE + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def status(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "get_state", {}) + value = getattr(value, "state", value) + state = TrajectoryState(value) + except (TypeError, ValueError, KeyError): + return Outcome.UNKNOWN + except Exception: + return Outcome.UNKNOWN + return { + TrajectoryState.IDLE: Outcome.INACTIVE, + TrajectoryState.EXECUTING: Outcome.RUNNING, + TrajectoryState.COMPLETED: Outcome.COMPLETED, + TrajectoryState.ABORTED: Outcome.CANCELLED, + TrajectoryState.FAULT: Outcome.FAILED, + }[state] + + def reset(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "reset", {}) + return ( + Outcome.INACTIVE + if value is True + else Outcome.RUNNING + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + try: + value = self._client.set_gripper_position(hardware_id, position) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def get_gripper_position(self, hardware_id: str) -> float | None: + try: + value = self._client.get_gripper_position(hardware_id) + return float(value) if isinstance(value, (int, float)) else None + except Exception: + return None + + def stop(self) -> None: + stop = getattr(self._client, "stop_rpc_client", None) + if callable(stop): + stop() diff --git a/dimos/manipulation/execution_models.py b/dimos/manipulation/execution_models.py new file mode 100644 index 0000000000..254e4cf31b --- /dev/null +++ b/dimos/manipulation/execution_models.py @@ -0,0 +1,188 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Pure public records and protocol models for execution runtime state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Generic, Protocol, TypeVar + +from dimos.manipulation.execution_topology import ExecutionPlan, PreparedPlan, TaskEntry + + +class LifecycleState(str, Enum): + IDLE = "IDLE" + PLANNING = "PLANNING" + READY = "READY" + DISPATCHING = "DISPATCHING" + RUNNING = "RUNNING" + CANCELLING = "CANCELLING" + FAULT = "FAULT" + + +class TaskActivity(str, Enum): + NOT_STARTED = "not_started" + EXECUTE_UNRESOLVED = "execute_unresolved" + ACTIVE = "active" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REMOTE_FAULT = "remote_fault" + UNKNOWN = "unknown" + + +class ActionMethod(str, Enum): + EXECUTE = "execute" + CANCEL = "cancel" + STATUS = "get_state" + RESET = "reset" + GRIPPER_SET = "set_gripper_position" + GRIPPER_GET = "get_gripper_position" + + +class Outcome(str, Enum): + ACCEPTED = "accepted" + RUNNING = "running" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REJECTED = "rejected" + FAILED = "failed" + UNKNOWN = "unknown" + + +class ShutdownState(str, Enum): + OPEN = "open" + CLOSING = "closing" + CLOSED = "closed" + + +@dataclass(frozen=True) +class OperationHandle: + plan_id: str + operation_id: str + attempt_id: str + + +ExecutionHandle = OperationHandle + + +@dataclass(frozen=True) +class ActionRecord: + action_id: str + method: ActionMethod + started: float + deadline: float + deadline_reported: bool = False + reset_id: str | None = None + + +@dataclass(frozen=True) +class TaskRecord: + task_id: str + task_name: str + entry: TaskEntry + activity: TaskActivity + action: ActionRecord | None = None + cancel_required: bool = False + reset_required: bool = False + + +PlanInput = ExecutionPlan | PreparedPlan + + +@dataclass(frozen=True) +class Operation: + handle: OperationHandle + plan: PlanInput + tasks: tuple[TaskRecord, ...] + next_index: int = 0 + cancel_requested: bool = False + uncertain: bool = False + rejected: bool = False + failed: bool = False + diagnostic: str = "" + + +@dataclass(frozen=True) +class ExecutionResult: + handle: OperationHandle + outcome: Outcome + diagnostic: str = "" + + +@dataclass(frozen=True) +class ResetHandle: + reset_id: str + + +@dataclass(frozen=True) +class ResetResult: + handle: ResetHandle + success: bool + diagnostic: str = "" + + +@dataclass(frozen=True) +class ShutdownResult: + success: bool + diagnostic: str = "" + + +@dataclass +class RuntimeContext: + state: LifecycleState = LifecycleState.IDLE + ready_plan: PlanInput | None = None + ready_plan_id: str | None = None + planning_token: str | None = None + active: Operation | None = None + fault: str | None = None + diagnostic: str | None = None + shutdown: ShutdownState = ShutdownState.OPEN + reset_handle: ResetHandle | None = None + dispatch_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + terminal_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + reset_results: tuple[tuple[ResetHandle, ResetResult], ...] = () + reset_proven_inactive: frozenset[str] = frozenset() + reset_completed_tasks: frozenset[str] = frozenset() + reset_deadline: float | None = None + shutdown_result: ShutdownResult | None = None + shutdown_deadline: float | None = None + physical_operation_deadline: float | None = None + revision: int = 0 + + +@dataclass(frozen=True) +class RuntimeSnapshot: + state: LifecycleState + ready_plan: PlanInput | None + ready_plan_id: str | None + planning_token: str | None + operation: Operation | None + fault: str | None + diagnostic: str | None + shutdown: ShutdownState + shutdown_result: ShutdownResult | None + revision: int + + +T = TypeVar("T") + + +@dataclass(frozen=True) +class CommandResult(Generic[T]): + accepted: bool + value: T | None = None + diagnostic: str = "" + snapshot: RuntimeSnapshot | None = None + + +class CoordinatorGateway(Protocol): + def execute(self, task_name: str, request: Any) -> Outcome: ... + def cancel(self, task_name: str) -> Outcome: ... + def status(self, task_name: str) -> Outcome: ... + def reset(self, task_name: str) -> Outcome: ... + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: ... + def get_gripper_position(self, hardware_id: str) -> float | None: ... + def stop(self) -> None: ... diff --git a/dimos/manipulation/execution_policy.py b/dimos/manipulation/execution_policy.py new file mode 100644 index 0000000000..4d08b79b7f --- /dev/null +++ b/dimos/manipulation/execution_policy.py @@ -0,0 +1,134 @@ +"""Pure policy decisions for execution-runtime reconciliation.""" + +from dataclasses import dataclass + +from dimos.manipulation.execution_models import ( + ActionMethod, + Outcome, + TaskActivity, +) + + +@dataclass(frozen=True) +class ResetCompletionDecision: + """Passive result of reconciling one reset-related coordinator call.""" + + activity: TaskActivity | None = None + cancel_required: bool | None = None + reset_required: bool | None = None + prove_inactive: bool = False + complete_reset: bool = False + advance_reset: bool = False + reset_success: bool | None = None + request_cancel: bool = False + emergency_cancel: bool = False + diagnostic: str | None = None + + +@dataclass(frozen=True) +class CompletionDecision: + """Pure normalized decision for a non-reset coordinator completion.""" + + activity: TaskActivity | None = None + fault: bool = False + rejected: bool = False + accepted: bool = False + diagnostic: str | None = None + + +_SAFE = frozenset((TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED)) + + +def reconcile_execute_completion(outcome: Outcome) -> CompletionDecision: + if outcome == Outcome.ACCEPTED: + return CompletionDecision(activity=TaskActivity.ACTIVE, accepted=True) + if outcome == Outcome.REJECTED: + return CompletionDecision(activity=TaskActivity.INACTIVE, rejected=True) + return CompletionDecision(fault=True, diagnostic=f"execute outcome={outcome.value}") + + +def reconcile_cancel_completion(outcome: Outcome) -> CompletionDecision: + if outcome in (Outcome.CANCELLED, Outcome.INACTIVE, Outcome.ACCEPTED): + return CompletionDecision(activity=TaskActivity.CANCELLED) + return CompletionDecision( + fault=True, + diagnostic=( + "cancellation is uncertain" + if outcome == Outcome.UNKNOWN + else "cancellation failed" + if outcome == Outcome.FAILED + else "malformed cancel outcome" + ), + ) + + +def reconcile_status_completion(outcome: Outcome) -> CompletionDecision: + activity = { + Outcome.COMPLETED: TaskActivity.COMPLETED, + Outcome.INACTIVE: TaskActivity.INACTIVE, + Outcome.CANCELLED: TaskActivity.CANCELLED, + Outcome.ACCEPTED: TaskActivity.ACTIVE, + Outcome.RUNNING: TaskActivity.ACTIVE, + }.get(outcome) + if activity is not None: + return CompletionDecision(activity=activity) + return CompletionDecision(fault=True, diagnostic=f"status outcome={outcome.value}") + + +def reconcile_reset_completion( + *, + method: ActionMethod, + outcome: Outcome, + activity: TaskActivity, + gate_sealed: bool, + clock_failed: bool, +) -> ResetCompletionDecision: + """Classify a reset completion without mutating runtime state.""" + if method == ActionMethod.STATUS: + if outcome in (Outcome.RUNNING, Outcome.ACCEPTED): + return ResetCompletionDecision( + activity=TaskActivity.ACTIVE, + cancel_required=True, + request_cancel=True, + emergency_cancel=gate_sealed or clock_failed, + ) + if outcome in (Outcome.INACTIVE, Outcome.CANCELLED, Outcome.COMPLETED): + resolved = { + Outcome.INACTIVE: TaskActivity.INACTIVE, + Outcome.CANCELLED: TaskActivity.CANCELLED, + Outcome.COMPLETED: TaskActivity.COMPLETED, + }[outcome] + return ResetCompletionDecision( + activity=resolved, + cancel_required=False, + prove_inactive=True, + advance_reset=True, + ) + safe = activity in _SAFE + return ResetCompletionDecision( + activity=activity if safe else TaskActivity.UNKNOWN, + cancel_required=not safe, + request_cancel=not safe, + emergency_cancel=not safe and (gate_sealed or clock_failed), + reset_success=False, + diagnostic="status reconciliation is unsafe", + ) + if method == ActionMethod.CANCEL: + if outcome in (Outcome.CANCELLED, Outcome.INACTIVE, Outcome.ACCEPTED): + return ResetCompletionDecision( + activity=TaskActivity.ACTIVE, + cancel_required=False, + advance_reset=True, + ) + return ResetCompletionDecision( + cancel_required=False, + reset_success=False, + diagnostic=None, + ) + if method == ActionMethod.RESET and outcome == Outcome.INACTIVE: + return ResetCompletionDecision( + reset_required=False, + complete_reset=True, + advance_reset=True, + ) + return ResetCompletionDecision(reset_success=False) diff --git a/dimos/manipulation/execution_runtime.py b/dimos/manipulation/execution_runtime.py index 33d35719cd..802a6709e4 100644 --- a/dimos/manipulation/execution_runtime.py +++ b/dimos/manipulation/execution_runtime.py @@ -1,590 +1,92 @@ # Copyright 2025-2026 Dimensional Inc. # Licensed under the Apache License, Version 2.0 (the "License"). -"""Owner-thread manipulation execution runtime. +"""Small, fail-closed owner-thread execution runtime. -Coordinator calls are effects. They never mutate runtime state; they only -post correlated events back to the owner. +The owner is the only thread which changes :class:`RuntimeContext`. Gateway +calls are submitted to a fixed executor and return as owner events; in +particular, an RPC can never hold the owner while it is in flight. """ from __future__ import annotations -from collections.abc import Callable, Iterable, Mapping -from dataclasses import dataclass, field, replace -from enum import Enum +from collections.abc import Callable +from dataclasses import dataclass, replace import queue import threading import time -from typing import Any, Generic, Protocol, TypeVar, cast +from typing import Any, TypeVar, cast import uuid -from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.execution_auxiliary import AuxiliaryCallBook, AuxiliaryTicket +from dimos.manipulation.execution_clock import InvalidMonotonicClock, ValidatedMonotonicClock +from dimos.manipulation.execution_effects import ( + AuxiliaryDone, + EffectDone, + ExecutionEffectRunner, + StopDone, +) +from dimos.manipulation.execution_gateway import ( + ControlCoordinatorGateway as ControlCoordinatorGateway, +) +from dimos.manipulation.execution_models import ( + ActionMethod as ActionMethod, + ActionRecord as ActionRecord, + CommandResult as CommandResult, + CoordinatorGateway as CoordinatorGateway, + ExecutionHandle as ExecutionHandle, + ExecutionResult as ExecutionResult, + LifecycleState as LifecycleState, + Operation as Operation, + OperationHandle as OperationHandle, + Outcome as Outcome, + PlanInput as PlanInput, + ResetHandle as ResetHandle, + ResetResult as ResetResult, + RuntimeContext as RuntimeContext, + RuntimeSnapshot as RuntimeSnapshot, + ShutdownResult as ShutdownResult, + ShutdownState as ShutdownState, + TaskActivity as TaskActivity, + TaskRecord as TaskRecord, +) +from dimos.manipulation.execution_policy import ( + reconcile_cancel_completion, + reconcile_execute_completion, + reconcile_reset_completion, + reconcile_status_completion, +) +from dimos.manipulation.execution_topology import ( + ExecutionPlan as ExecutionPlan, + ExecutionTopology as ExecutionTopology, + PreparedPlan as PreparedPlan, + TaskEntry as TaskEntry, + materialize_prepared_plan as materialize_prepared_plan, + prepare_execution_plan as prepare_execution_plan, + prepare_generated_plan as prepare_generated_plan, + prepare_plan as prepare_plan, +) from dimos.manipulation.planning.spec.models import GeneratedPlan -from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +T = TypeVar("T") -class LifecycleState(str, Enum): - IDLE = "IDLE" - PLANNING = "PLANNING" - READY = "READY" - DISPATCHING = "DISPATCHING" - RUNNING = "RUNNING" - CANCELLING = "CANCELLING" - FAULT = "FAULT" - - -class TaskActivity(str, Enum): - NOT_STARTED = "not_started" - EXECUTE_UNRESOLVED = "execute_unresolved" - ACTIVE = "active" - COMPLETED = "completed" - CANCELLED = "cancelled" - INACTIVE = "inactive" - REMOTE_FAULT = "remote_fault" - UNKNOWN = "unknown" - - -class ActionMethod(str, Enum): - EXECUTE = "execute" - CANCEL = "cancel" - STATUS = "get_state" - RESET = "reset" - GRIPPER_SET = "set_gripper_position" - GRIPPER_GET = "get_gripper_position" - - -class Outcome(str, Enum): - ACCEPTED = "accepted" - RUNNING = "running" - COMPLETED = "completed" - CANCELLED = "cancelled" - INACTIVE = "inactive" - REJECTED = "rejected" - FAILED = "failed" - UNKNOWN = "unknown" - - -class ShutdownState(str, Enum): - OPEN = "open" - CLOSING = "closing" - CLOSED = "closed" - - -class _GatewayGateState(str, Enum): - OPEN = "open" - CANCEL_ONLY = "cancel_only" - SEALED = "sealed" - STOPPED = "stopped" - - -@dataclass(frozen=True) -class ExecutionTopology: - """Canonical group -> robot -> coordinator-task routing.""" - - routes: tuple[tuple[str, str, str], ...] - robot_configs: tuple[RobotModelConfig, ...] = () - inverse_joint_mappings: tuple[tuple[str, tuple[tuple[str, str], ...]], ...] = () - - @classmethod - def from_robot_configs( - cls, - robot_configs: Iterable[RobotModelConfig], - group_routes: Mapping[str, tuple[str, str]] | Iterable[tuple[str, str, str]] | None = None, - group_resolver: Callable[[str], tuple[str, str]] | None = None, - ) -> ExecutionTopology: - """Build canonical routes while retaining the supplied config references.""" - configs = tuple(robot_configs) - by_robot: dict[str, RobotModelConfig] = {} - for config in configs: - if config.name in by_robot: - raise ValueError(f"duplicate robot config: {config.name}") - by_robot[config.name] = config - explicit: dict[str, tuple[str, str]] = {} - if isinstance(group_routes, Mapping): - for group, route in group_routes.items(): - if not isinstance(group, str) or not isinstance(route, tuple) or len(route) != 2: - raise ValueError("invalid group route mapping") - explicit[group] = (str(route[0]), str(route[1])) - elif group_routes is not None: - for group, robot, task in group_routes: - if group in explicit: - raise ValueError(f"duplicate group route: {group}") - explicit[group] = (robot, task) - routes: list[tuple[str, str, str]] = [] - inverse: list[tuple[str, tuple[tuple[str, str], ...]]] = [] - for config in configs: - local_names = set(config.joint_names) - local_to_coordinator: dict[str, str] = {} - for coordinator_name, local_name in config.joint_name_mapping.items(): - if local_name not in local_names: - raise ValueError( - f"mapped local joint is not configured: {config.name}/{local_name}" - ) - previous = local_to_coordinator.get(local_name) - if previous is not None and previous != coordinator_name: - raise ValueError( - f"duplicate local joint mapping for {config.name}/{local_name}" - ) - local_to_coordinator[local_name] = coordinator_name - coordinator_names: set[str] = set() - for local_name in config.joint_names: - coordinator_name = local_to_coordinator.get(local_name, local_name) - if coordinator_name in coordinator_names: - raise ValueError( - f"duplicate coordinator joint mapping for {config.name}/{coordinator_name}" - ) - coordinator_names.add(coordinator_name) - inverse.append( - ( - config.name, - tuple( - (local, local_to_coordinator.get(local, local)) - for local in config.joint_names - ), - ) - ) - for definition in config.planning_groups: - if not set(definition.joint_names).issubset(local_names): - raise ValueError( - f"planning group joints are not configured: {config.name}/{definition.name}" - ) - group_id = f"{config.name}/{definition.name}" - if group_resolver is not None: - route_robot, route_task = group_resolver(group_id) - else: - if not config.coordinator_task_name and group_id not in explicit: - continue - route_robot, route_task = explicit.get( - group_id, (config.name, config.coordinator_task_name or "") - ) - if route_robot != config.name or route_task != config.coordinator_task_name: - raise ValueError(f"route/config mismatch for {group_id}") - routes.append((group_id, route_robot, route_task)) - if explicit and set(explicit) != {group for group, _, _ in routes}: - raise ValueError("route set does not match configured planning groups") - return cls(tuple(routes), configs, tuple(inverse)) - - from_configs = from_robot_configs - from_robot_model_configs = from_robot_configs - - def config_for_robot(self, robot_name: str) -> RobotModelConfig: - for config in self.robot_configs: - if config.name == robot_name: - return config - raise KeyError(f"unknown robot: {robot_name}") - - def coordinator_joint_name(self, robot_name: str, local_name: str) -> str: - for name, mappings in self.inverse_joint_mappings: - if name == robot_name: - for local, coordinator in mappings: - if local == local_name: - return coordinator - break - return local_name - - def route_for_group(self, group_id: str) -> tuple[str, str]: - routes = {group: (robot, task) for group, robot, task in self.routes} - try: - return routes[group_id] - except KeyError as exc: - raise ValueError(f"missing route for planning group {group_id}") from exc - - def affected_robots(self, group_ids: Iterable[str]) -> tuple[str, ...]: - robots: list[str] = [] - for group_id in group_ids: - robot_name, _task_name = self.route_for_group(group_id) - if robot_name in robots: - raise ValueError(f"multiple selected groups resolve to robot {robot_name}") - robots.append(robot_name) - return tuple(robots) - @property - def groups(self) -> tuple[str, ...]: - return tuple(group for group, _, _ in self.routes) +class _Call: + def __init__(self, fn: Callable[[], Any]) -> None: + self.fn = fn + self.answer: queue.Queue[Any] = queue.Queue(1) @dataclass(frozen=True) -class TaskEntry: - planning_group: str +class _PendingAction: + handle: OperationHandle task_name: str - request: Any - robot_name: str | None = None - - -@dataclass(frozen=True) -class ExecutionPlan: - generated_plan: Any - entries: tuple[TaskEntry, ...] - - -@dataclass(frozen=True) -class PreparedPlan: - generated_plan: GeneratedPlan - entries: tuple[TaskEntry, ...] - topology: ExecutionTopology | None = None - _canonical: bool = field(default=False, repr=False, compare=False) - - -def prepare_generated_plan( - generated_plan: GeneratedPlan, topology: ExecutionTopology -) -> PreparedPlan: - """Purely bind a generated multi-robot plan to canonical coordinator tasks.""" - if not isinstance(generated_plan, GeneratedPlan): - raise TypeError("expected GeneratedPlan") - group_ids = tuple(generated_plan.group_ids) - if not group_ids or len(set(group_ids)) != len(group_ids): - raise ValueError("generated plan must contain unique planning groups") - route_by_group = {group: (robot, task) for group, robot, task in topology.routes} - if len(route_by_group) != len(topology.routes): - raise ValueError("topology contains duplicate group routes") - selected_routes: list[tuple[str, str, str, RobotModelConfig, tuple[str, ...]]] = [] - affected: set[str] = set() - for group_id in group_ids: - route = route_by_group.get(group_id) - if route is None: - robot_name = group_id.split("/", 1)[0] - try: - if not topology.config_for_robot(robot_name).coordinator_task_name: - raise ValueError(f"no coordinator task for selected robot {robot_name}") - except KeyError: - pass - raise ValueError(f"missing route for planning group {group_id}") - robot_name, task_name = route - if not task_name: - raise ValueError(f"missing coordinator task for {group_id}") - if robot_name in affected: - raise ValueError(f"multiple selected groups resolve to robot {robot_name}") - config = topology.config_for_robot(robot_name) - prefix, _, definition_name = group_id.partition("/") - if prefix != robot_name or not definition_name: - raise ValueError(f"invalid canonical planning group {group_id}") - definitions = [ - definition - for definition in config.planning_groups - if definition.name == definition_name - ] - if len(definitions) != 1: - raise ValueError(f"missing config group for {group_id}") - expected_local = tuple(definitions[0].joint_names) - if len(set(expected_local)) != len(expected_local): - raise ValueError(f"duplicate intended joints for {group_id}") - affected.add(robot_name) - selected_routes.append((group_id, robot_name, task_name, config, expected_local)) - expected_by_robot = {robot: set(expected) for _, robot, _, _, expected in selected_routes} - - trajectory = generated_plan.trajectory - names = tuple(trajectory.joint_names) - if not names: - raise ValueError("generated trajectory has no joint names") - if len(set(names)) != len(names): - raise ValueError("duplicate global trajectory joint assignment") - columns: dict[str, dict[str, int]] = {robot: {} for robot in affected} - known_robots = {config.name for config in topology.robot_configs} - for index, global_name in enumerate(names): - if global_name.count("/") != 1: - raise ValueError(f"malformed global joint name: {global_name}") - robot_name, local_name = global_name.split("/", 1) - if not robot_name or not local_name or robot_name not in known_robots: - raise ValueError(f"unknown global joint assignment: {global_name}") - if robot_name in columns: - if local_name not in expected_by_robot[robot_name]: - raise ValueError(f"additional joint column for selected robot: {global_name}") - if local_name in columns[robot_name]: - raise ValueError(f"duplicate robot joint assignment: {global_name}") - columns[robot_name][local_name] = index - points = tuple(trajectory.points) - for point in points: - if len(point.positions) != len(names): - raise ValueError("trajectory point dimension does not match joint names") - if point.velocities and len(point.velocities) != len(names): - raise ValueError("trajectory velocity dimension does not match joint names") - - entries: list[TaskEntry] = [] - for group_id, robot_name, task_name, _config, expected_local in selected_routes: - robot_columns = columns[robot_name] - missing = [local for local in expected_local if local not in robot_columns] - if missing: - raise ValueError(f"missing robot joints for {robot_name}: {missing}") - intended = set(expected_local) - local_order = tuple(local for local in robot_columns if local in intended) - coordinator_names = tuple( - topology.coordinator_joint_name(robot_name, local) for local in local_order - ) - if len(set(coordinator_names)) != len(coordinator_names): - raise ValueError(f"duplicate coordinator joint names for {robot_name}") - local_trajectory = JointTrajectory( - joint_names=list(coordinator_names), - points=[ - TrajectoryPoint( - time_from_start=point.time_from_start, - positions=[point.positions[robot_columns[local]] for local in local_order], - velocities=( - [point.velocities[robot_columns[local]] for local in local_order] - if point.velocities - else None - ), - ) - for point in points - ], - ) - entries.append(TaskEntry(group_id, task_name, {"trajectory": local_trajectory}, robot_name)) - return PreparedPlan(generated_plan, tuple(entries), topology, True) - - -prepare_plan = prepare_generated_plan -prepare_execution_plan = prepare_generated_plan -materialize_prepared_plan = prepare_generated_plan - -PlanInput = ExecutionPlan | PreparedPlan - - -@dataclass(frozen=True) -class OperationHandle: - plan_id: str - operation_id: str - attempt_id: str - - -ExecutionHandle = OperationHandle - - -@dataclass(frozen=True) -class ActionRecord: - action_id: str + task: TaskRecord method: ActionMethod - started: float deadline: float - deadline_reported: bool = False - reset_id: str | None = None - - -@dataclass(frozen=True) -class TaskRecord: - task_id: str - task_name: str - entry: TaskEntry - activity: TaskActivity - action: ActionRecord | None = None - cancel_required: bool = False - reset_required: bool = False - - -@dataclass(frozen=True) -class Operation: - handle: OperationHandle - plan: PlanInput - tasks: tuple[TaskRecord, ...] - next_index: int = 0 - cancel_requested: bool = False - cleanup_required: bool = False - uncertain: bool = False - rejected: bool = False - failed: bool = False - diagnostic: str = "" - - -@dataclass(frozen=True) -class ExecutionResult: - handle: OperationHandle - outcome: Outcome - diagnostic: str = "" - - -@dataclass(frozen=True) -class ResetHandle: - reset_id: str - - -@dataclass(frozen=True) -class ResetResult: - handle: ResetHandle - success: bool - diagnostic: str = "" - - -@dataclass(frozen=True) -class ShutdownResult: - success: bool - diagnostic: str = "" - - -@dataclass(frozen=True) -class RuntimeContext: - state: LifecycleState = LifecycleState.IDLE - ready_plan: PlanInput | None = None - ready_plan_id: str | None = None - planning_token: str | None = None - active: Operation | None = None - fault: str | None = None - diagnostic: str | None = None - shutdown: ShutdownState = ShutdownState.OPEN - reset_handle: ResetHandle | None = None - reset_result: ResetResult | None = None - retired_plan_ids: frozenset[str] = frozenset() - dispatch_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () - terminal_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () - reset_results: tuple[tuple[ResetHandle, ResetResult], ...] = () - reset_proven_inactive: frozenset[str] = frozenset() - reset_completed_tasks: frozenset[str] = frozenset() - reset_deadline: float | None = None - shutdown_result: ShutdownResult | None = None - shutdown_deadline: float | None = None - revision: int = 0 - - -@dataclass(frozen=True) -class RuntimeSnapshot: - state: LifecycleState - ready_plan: PlanInput | None - ready_plan_id: str | None - planning_token: str | None - operation: Operation | None - fault: str | None - diagnostic: str | None - shutdown: ShutdownState - shutdown_result: ShutdownResult | None - revision: int - - -T = TypeVar("T") - - -@dataclass(frozen=True) -class CommandResult(Generic[T]): - accepted: bool - value: T | None = None - diagnostic: str = "" - snapshot: RuntimeSnapshot | None = None - - -class CoordinatorGateway(Protocol): - def execute(self, task_name: str, request: Any) -> Outcome: ... - def cancel(self, task_name: str) -> Outcome: ... - def status(self, task_name: str) -> Outcome: ... - def reset(self, task_name: str) -> Outcome: ... - def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: ... - def get_gripper_position(self, hardware_id: str) -> float | None: ... - def stop(self) -> None: ... - - -class ControlCoordinatorGateway: - """Adapter for ControlCoordinator.task_invoke(task, method, kwargs).""" - - def __init__(self, client: Any) -> None: - self._client = client - - def _invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: - return self._client.task_invoke(task, method, kwargs) - - def execute(self, task_name: str, request: Any) -> Outcome: - try: - value = self._invoke( - task_name, - "execute", - request if isinstance(request, dict) else {"trajectory": request}, - ) - return ( - Outcome.ACCEPTED - if value is True - else Outcome.REJECTED - if value is False - else Outcome.UNKNOWN - ) - except Exception: - return Outcome.UNKNOWN - - def cancel(self, task_name: str) -> Outcome: - try: - value = self._invoke(task_name, "cancel", {}) - return ( - Outcome.CANCELLED - if value is True - else Outcome.INACTIVE - if value is False - else Outcome.UNKNOWN - ) - except Exception: - return Outcome.UNKNOWN - - def status(self, task_name: str) -> Outcome: - try: - value = self._invoke(task_name, "get_state", {}) - value = getattr(value, "state", value) - state = TrajectoryState(value) - except (TypeError, ValueError, KeyError): - return Outcome.UNKNOWN - except Exception: - return Outcome.UNKNOWN - return { - TrajectoryState.IDLE: Outcome.INACTIVE, - TrajectoryState.EXECUTING: Outcome.RUNNING, - TrajectoryState.COMPLETED: Outcome.COMPLETED, - TrajectoryState.ABORTED: Outcome.CANCELLED, - TrajectoryState.FAULT: Outcome.FAILED, - }[state] - - def reset(self, task_name: str) -> Outcome: - try: - value = self._invoke(task_name, "reset", {}) - return ( - Outcome.INACTIVE - if value is True - else Outcome.RUNNING - if value is False - else Outcome.UNKNOWN - ) - except Exception: - return Outcome.UNKNOWN - - def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: - try: - value = self._client.set_gripper_position(hardware_id, position) - return ( - Outcome.ACCEPTED - if value is True - else Outcome.REJECTED - if value is False - else Outcome.UNKNOWN - ) - except Exception: - return Outcome.UNKNOWN - - def get_gripper_position(self, hardware_id: str) -> float | None: - try: - value = self._client.get_gripper_position(hardware_id) - return float(value) if isinstance(value, (int, float)) else None - except Exception: - return None - - def stop(self) -> None: - stop = getattr(self._client, "stop_rpc_client", None) - if callable(stop): - stop() - - -@dataclass(frozen=True) -class _Event: - kind: str - operation_id: str | None = None - attempt_id: str | None = None - task_id: str | None = None - action_id: str | None = None - method: ActionMethod | None = None - outcome: Outcome | None = None - diagnostic: str = "" - payload: Any = None - reset_id: str | None = None - - -@dataclass -class _Call: - function: Callable[[], Any] - answer: queue.Queue[Any] = field(default_factory=lambda: queue.Queue(1)) class ExecutionRuntime: - """All context commits and effect launches happen on one owner thread.""" + """Serialized plan executor with a sticky, fail-closed fault state.""" def __init__( self, @@ -593,163 +95,40 @@ def __init__( topology: ExecutionTopology | None = None, plan_validator: Callable[[PlanInput], bool] | None = None, action_timeout: float = 1.0, + physical_operation_timeout: float = 60.0, poll_interval: float = 0.1, + monotonic_clock: Callable[[], float] = time.monotonic, ) -> None: self._gateway = gateway_factory() self._topology = topology - self._validate_plan = plan_validator or self._valid_plan - self._timeout = action_timeout - self._poll_interval = poll_interval + self._validator = plan_validator + self._action_timeout = max(0.0, action_timeout) + self._physical_timeout = max(0.0, physical_operation_timeout) + self._poll_interval = max(0.001, poll_interval) + self._clock = ValidatedMonotonicClock(monotonic_clock) self._context = RuntimeContext() - self._effects: list[tuple[Operation, TaskRecord, ActionMethod]] = [] - self._context_lock = threading.Lock() - self._events: queue.Queue[_Call | None] = queue.Queue() - self._closed = False - self._stop = threading.Event() - self._condition = threading.Condition(self._context_lock) - self._owner = threading.Thread( - target=self._loop, daemon=True, name="manipulation-execution-owner" - ) - self._owner.start() - self._poller = threading.Thread( - target=self._poll_loop, daemon=True, name="manipulation-execution-poller" + self._lock = threading.RLock() + self._condition = threading.Condition(self._lock) + self._events: queue.Queue[_Call | EffectDone | StopDone | AuxiliaryDone | None] = ( + queue.Queue() ) - self._poller.start() + self._effects = ExecutionEffectRunner() + self._closed = False + self._gate = "open" + self._pending: dict[str, _PendingAction] = {} + self._fault_terminal_results: dict[OperationHandle, str] = {} + self._next_poll: float | None = None self._shutdown_lock = threading.Lock() - self._auxiliary_lock = threading.Lock() - self._auxiliary: dict[threading.Thread, threading.Event] = {} - self._closing_requested = False - self._gateway_condition = threading.Condition() - self._gateway_gate = _GatewayGateState.OPEN - self._gateway_close_started = False - self._gateway_close_done = False - self._gateway_close_error: str | None = None - self._rpc_inflight = 0 - self._admitted_actions: set[str] = set() - self._started_actions: set[str] = set() - - def _start_auxiliary( - self, - target: Callable[..., None], - *args: Any, - name: str, - ) -> None: - """Start a tracked effect thread with cooperative cancellation.""" - cancel = threading.Event() - - def run() -> None: - try: - target(cancel, *args) - finally: - with self._auxiliary_lock: - self._auxiliary.pop(threading.current_thread(), None) - - thread = threading.Thread(target=run, name=name, daemon=True) - with self._auxiliary_lock: - self._auxiliary[thread] = cancel - thread.start() - - def _stop_auxiliary(self, deadline: float) -> bool: - """Cancel and join tracked effects until the caller's deadline.""" - with self._auxiliary_lock: - pending = tuple(self._auxiliary.items()) - for _thread, cancel in pending: - cancel.set() - for thread, _cancel in pending: - thread.join(max(0.0, deadline - time.monotonic())) - with self._auxiliary_lock: - return not self._auxiliary - - def _valid_plan(self, plan: PlanInput) -> bool: - if not isinstance(plan, (ExecutionPlan, PreparedPlan)) or not plan.entries: - return False - if isinstance(plan, ExecutionPlan) and isinstance(plan.generated_plan, GeneratedPlan): - return False - groups = tuple(getattr(plan.generated_plan, "group_ids", ())) - represented = tuple(entry.planning_group for entry in plan.entries) - if len(set(represented)) != len(represented) or len( - set(entry.task_name for entry in plan.entries) - ) != len(plan.entries): - return False - if isinstance(plan, PreparedPlan) and ( - not plan._canonical - or not isinstance(plan.generated_plan, GeneratedPlan) - or plan.topology is None - ): - return False - if isinstance(plan, PreparedPlan) and groups != represented: - return False - if not isinstance(plan, PreparedPlan) and groups and set(groups) != set(represented): - return False - topology = plan.topology if isinstance(plan, PreparedPlan) else self._topology - if topology is not None: - routes = {group: (robot, task) for group, robot, task in topology.routes} - selected_groups = groups or represented - if len(routes) != len(topology.routes) or any( - group not in routes for group in selected_groups - ): - return False - selected = tuple(routes[group] for group in selected_groups) - if tuple(entry.planning_group for entry in plan.entries) != selected_groups: - return False - if any( - ( - route[1] != entry.task_name - or (entry.robot_name is not None and route[0] != entry.robot_name) - ) - for entry, route in zip(plan.entries, selected, strict=True) - ): - return False - return True - - def snapshot(self) -> RuntimeSnapshot: - with self._context_lock: - c = self._context - return RuntimeSnapshot( - c.state, - c.ready_plan, - c.ready_plan_id, - c.planning_token, - c.active, - c.fault, - c.diagnostic, - c.shutdown, - c.shutdown_result, - c.revision, - ) - - def _loop(self) -> None: - while True: - call = self._events.get() - if call is None: - return - try: - value = call.function() - self._process_effects() - call.answer.put(value) - except BaseException as exc: - call.answer.put(exc) - - def _submit(self, fn: Callable[[], T], *, allow_closing: bool = False) -> T: - with self._context_lock: - reject = self._closed or ( - not allow_closing - and (self._closing_requested or self._context.shutdown != ShutdownState.OPEN) - ) - if reject: - diagnostic = "runtime is closed" if self._closed else "runtime is closing" - return cast( - "T", - CommandResult(False, diagnostic=diagnostic, snapshot=self._snapshot_unlocked()), - ) - call = _Call(fn) - self._events.put(call) - result = call.answer.get() - if isinstance(result, BaseException): - raise result - return cast("T", result) + self._shutdown_initiated = False + self._shutdown_stop_started = False + self._shutdown_stop_finished = False + self._shutdown_drain_started = False + self._clock_invalid = False + self._auxiliary = AuxiliaryCallBook() + self._owner = threading.Thread(target=self._run, daemon=True, name="execution-owner") + self._owner.start() - def _snapshot_unlocked(self) -> RuntimeSnapshot: + def _snapshot(self) -> RuntimeSnapshot: c = self._context return RuntimeSnapshot( c.state, @@ -764,1086 +143,1190 @@ def _snapshot_unlocked(self) -> RuntimeSnapshot: c.revision, ) - def _commit(self, context: RuntimeContext) -> RuntimeContext: - with self._condition: - self._context = replace(context, revision=context.revision + 1) + def snapshot(self) -> RuntimeSnapshot: + with self._lock: + return self._snapshot() + + def _commit(self, **kwargs: Any) -> None: + with self._lock: + self._context = replace(self._context, **kwargs, revision=self._context.revision + 1) self._condition.notify_all() - return self._context - def _record_dispatch_result_if_absent(self, result: ExecutionResult) -> None: - """Record the first dispatch outcome; later cleanup never replaces it.""" - if dict(self._context.dispatch_results).get(result.handle) is None: - self._commit( - replace( - self._context, - dispatch_results=(*self._context.dispatch_results, (result.handle, result))[ - -16: - ], + def _submit(self, fn: Callable[[], T], *, closing: bool = False) -> T: + with self._lock: + if self._closed or ( + self._shutdown_drain_started + or (self._shutdown_initiated and not self._owner.is_alive()) + or ( + not closing + and (self._context.shutdown != ShutdownState.OPEN or self._gate != "open") ) - ) - - def _reset_pending(self) -> bool: - handle = self._context.reset_handle - return ( - self._context.state == LifecycleState.FAULT - and handle is not None - and handle not in dict(self._context.reset_results) - ) + ): + return cast( + "T", + CommandResult( + False, diagnostic="runtime is closing", snapshot=self._snapshot() + ), + ) + call = _Call(fn) + self._events.put(call) + result = call.answer.get() + if isinstance(result, BaseException): + raise result + return cast("T", result) - def _finish_reset(self, success: bool, diagnostic: str = "") -> None: - handle = self._context.reset_handle - if handle is None or handle in dict(self._context.reset_results): - return - result = ResetResult(handle, success, diagnostic) - reset_results = (*self._context.reset_results, (handle, result))[-16:] - self._commit( - replace( - self._context, - state=LifecycleState.IDLE if success else LifecycleState.FAULT, - active=None if success else self._context.active, - fault=None if success else self._context.fault, - diagnostic=None if success else diagnostic, - reset_results=reset_results, - # A failed attempt is retained in reset_results for its - # caller, but must not remain the active reconciliation gate: - # a later reset command must be able to start a fresh attempt. - reset_handle=None, - reset_proven_inactive=frozenset(), - reset_completed_tasks=frozenset(), - reset_deadline=None, + def _run(self) -> None: + while True: + timeout = self._wait_timeout() + try: + event = self._events.get(timeout=timeout) + except queue.Empty: + with self._lock: + self._timers() + continue + if event is None: + return + try: + with self._lock: + if isinstance(event, _Call): + self._timers() + event.answer.put(event.fn()) + elif isinstance(event, EffectDone): + self._done(event) + elif isinstance(event, AuxiliaryDone): + self._aux_done(event) + else: + self._stop_done(event) + self._timers() + except BaseException as exc: + if isinstance(event, _Call): + event.answer.put(exc) + + def _wait_timeout(self) -> float: + if self._clock_invalid: + return 3600.0 + deadlines = [ + action.deadline + for action in self._pending.values() + if action.task.action is None or not action.task.action.deadline_reported + ] + deadlines.extend(self._auxiliary.deadlines()) + if self._next_poll is not None: + deadlines.append(self._next_poll) + if self._context.reset_deadline is not None: + deadlines.append(self._context.reset_deadline) + if self._context.shutdown_deadline is not None: + deadlines.append(self._context.shutdown_deadline) + if self._context.physical_operation_deadline is not None: + deadlines.append(self._context.physical_operation_deadline) + if not deadlines: + return 3600.0 + now = self._clock_now() + if now is None: + return 3600.0 + return max(0.0, min(deadlines) - now) + + def _clock_failure_owner(self, diagnostic: str) -> None: + self._gate = "sealed" + self._shutdown_initiated = True + op = self._context.active + if op is not None: + potentially_active = { + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + } + op = replace( + op, + tasks=tuple( + replace( + task, + cancel_required=task.cancel_required or task.activity in potentially_active, + reset_required=(task.reset_required or task.activity in potentially_active), + ) + for task in op.tasks + ), + uncertain=True, + diagnostic=diagnostic, ) + self._commit( + state=LifecycleState.FAULT + if self._context.state == LifecycleState.FAULT + else LifecycleState.CANCELLING, + shutdown=ShutdownState.CLOSING, + shutdown_result=self._context.shutdown_result or ShutdownResult(False, diagnostic), + active=op, + diagnostic=diagnostic, ) + if op is not None: + self._schedule_cancels(op, emergency=True) + self._shutdown_progress() - def _finish_shutdown(self, success: bool, diagnostic: str = "") -> None: - if self._context.shutdown_result is not None: - return - self._commit(replace(self._context, shutdown_result=ShutdownResult(success, diagnostic))) + def _clock_now(self) -> float | None: + if self._clock_invalid: + return None + try: + return self._clock.now() + except InvalidMonotonicClock: + self._clock_invalid = True + if threading.current_thread() is self._owner: + self._clock_failure_owner("invalid monotonic clock") + else: + with self._lock: + if not self._shutdown_drain_started and self._owner.is_alive(): + self._events.put( + _Call(lambda: self._clock_failure_owner("invalid monotonic clock")) + ) + return None - def _maybe_finish_shutdown(self) -> None: + def _timers(self) -> None: + now = self._clock_now() + if now is None: + return + for action_id, item in tuple(self._pending.items()): + if item.deadline <= now: + handle, task_name, task, method = ( + item.handle, + item.task_name, + item.task, + item.method, + ) + op = self._context.active + report = False + if op is not None and op.handle == handle: + current = next((t for t in op.tasks if t.task_id == task.task_id), None) + if ( + current is not None + and current.action is not None + and not current.action.deadline_reported + ): + report = True + action = replace(current.action, deadline_reported=True) + tasks = tuple( + replace( + t, + action=action, + activity=TaskActivity.UNKNOWN, + cancel_required=True, + ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + self._commit(active=replace(op, tasks=tasks, uncertain=True)) + self._pending[action_id] = replace( + item, + task=replace(task, action=action), + ) + if report: + self._fault( + f"coordinator {method.value} deadline for {task_name}", + handle, + task, + ) + if self._auxiliary.expire(now, "coordinator auxiliary action deadline exceeded"): + self._condition.notify_all() + physical_deadline = self._context.physical_operation_deadline if ( - self._context.shutdown != ShutdownState.CLOSING - or self._context.shutdown_result is not None + physical_deadline is not None + and now >= physical_deadline + and self._context.active is not None + and self._context.state == LifecycleState.RUNNING ): + self._commit(physical_operation_deadline=None) + self._begin_cancellation("physical operation deadline exceeded") + if self._next_poll is not None and now >= self._next_poll: + self._next_poll = now + self._poll_interval + self._poll_active() + if self._context.reset_deadline is not None and now >= self._context.reset_deadline: + if self._context.reset_handle and self._context.reset_handle not in dict( + self._context.reset_results + ): + self._finish_reset(False, "reset deadline exceeded") + if self._context.shutdown_deadline is not None and now >= self._context.shutdown_deadline: + if self._context.shutdown_result is None: + self._finish_shutdown(False, "shutdown safety deadline exceeded") + + def _valid_plan(self, plan: PlanInput) -> bool: + if not isinstance(plan, (ExecutionPlan, PreparedPlan)) or not plan.entries: + return False + if isinstance(plan, ExecutionPlan) and isinstance(plan.generated_plan, GeneratedPlan): + return False + if isinstance(plan, PreparedPlan) and (not plan._canonical or plan.topology is None): + return False + groups = tuple(getattr(plan.generated_plan, "group_ids", ())) + names = tuple(e.planning_group for e in plan.entries) + if len(set(names)) != len(names) or len({e.task_name for e in plan.entries}) != len(names): + return False + if groups and groups != names: + return False + topology = plan.topology if isinstance(plan, PreparedPlan) else self._topology + if topology is not None: + routes = {g: (r, t) for g, r, t in topology.routes} + for entry in plan.entries: + route = routes.get(entry.planning_group) + if ( + route is None + or route[1] != entry.task_name + or (entry.robot_name and route[0] != entry.robot_name) + ): + return False + return True + + def _valid(self, plan: PlanInput) -> bool: + try: + return bool(self._validator(plan) if self._validator else self._valid_plan(plan)) + except Exception: + return False + + def _start_action( + self, + op: Operation, + task: TaskRecord, + method: ActionMethod, + *, + emergency: bool = False, + ) -> None: + if task.action is not None or (self._gate == "sealed" and not emergency): return - operation = self._context.active - if operation is None: - self._finish_shutdown(True) + if self._gate != "open" and not ( + method == ActionMethod.CANCEL + and task.cancel_required + and task.activity != TaskActivity.NOT_STARTED + and (emergency or self._gate == "cancel_only") + ): return - if any(task.action is not None for task in operation.tasks): + if method not in (ActionMethod.STATUS, ActionMethod.RESET) and task.activity in ( + TaskActivity.INACTIVE, + TaskActivity.COMPLETED, + TaskActivity.CANCELLED, + ): return - if any(task.activity == TaskActivity.REMOTE_FAULT for task in operation.tasks): - self._finish_shutdown(False, "shutdown safety unresolved: remote fault") + now = self._clock_now() + if now is None and method != ActionMethod.CANCEL: return - pending = next( - ( - task - for task in operation.tasks - if task.activity - in ( - TaskActivity.NOT_STARTED, - TaskActivity.EXECUTE_UNRESOLVED, - TaskActivity.ACTIVE, - TaskActivity.UNKNOWN, - ) - ), - None, + action = ActionRecord( + str(uuid.uuid4()), + method, + now or 0.0, + float("inf") if now is None else now + self._action_timeout, + ) + tasks = tuple( + replace( + t, + action=action, + activity=TaskActivity.EXECUTE_UNRESOLVED + if method == ActionMethod.EXECUTE + else t.activity, + ) + if t.task_id == task.task_id + else t + for t in op.tasks ) - if pending is not None: - self._schedule_action(operation, pending, ActionMethod.CANCEL) + op = replace(op, tasks=tasks) + self._commit(active=op) + self._pending[action.action_id] = _PendingAction( + handle=op.handle, + task_name=task.task_name, + task=next(t for t in tasks if t.task_id == task.task_id), + method=method, + deadline=action.deadline, + ) + fn: Callable[[], Any] + if method == ActionMethod.EXECUTE: + fn = lambda: self._gateway.execute(task.task_name, task.entry.request) # noqa: E731 + elif method == ActionMethod.CANCEL: + fn = lambda: self._gateway.cancel(task.task_name) # noqa: E731 + elif method == ActionMethod.STATUS: + fn = lambda: self._gateway.status(task.task_name) # noqa: E731 + else: + fn = lambda: self._gateway.reset(task.task_name) # noqa: E731 + self._effects.submit_action(action.action_id, fn, self._events.put) + + @dataclass(frozen=True) + class _RetiredAction: + operation: Operation + index: int + task: TaskRecord + method: ActionMethod + + def _retire_action(self, done: EffectDone) -> _RetiredAction | None: + item = self._pending.get(done.action_id) + active = self._context.active + if item is None or active is None: + return None + if active.handle != item.handle: + return None + index = next( + (i for i, t in enumerate(active.tasks) if t.task_id == item.task.task_id), None + ) + if index is None: + return None + current_task = active.tasks[index] + if current_task.action is None or current_task.action.action_id != done.action_id: + # The registered operation/task identity is sufficient to consume this + # physical completion, but it must not retire a newer action. + self._pending.pop(done.action_id, None) + return None + self._pending.pop(done.action_id, None) + task = replace(active.tasks[index], action=None) + operation = replace(active, tasks=(*active.tasks[:index], task, *active.tasks[index + 1 :])) + self._commit(active=operation) + return self._RetiredAction(operation, index, task, item.method) + + def _done(self, done: EffectDone) -> None: + retired = self._retire_action(done) + if retired is None: return - if all( - task.activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED) - for task in operation.tasks + op, index, task, method = ( + retired.operation, + retired.index, + retired.task, + retired.method, + ) + outcome = done.outcome if isinstance(done.outcome, Outcome) else Outcome.UNKNOWN + if self._context.reset_handle is not None and method in ( + ActionMethod.STATUS, + ActionMethod.CANCEL, + ActionMethod.RESET, ): - self._finish_shutdown(True) - - def _reconcile_reset(self, op: Operation | None, event: _Event | None = None) -> None: - if not self._reset_pending(): + self._done_reset(op, index, task, method, outcome) return - op = op or self._context.active - if op is None: - self._finish_reset(False, "fault has no operation to reconcile") + if method == ActionMethod.EXECUTE: + self._done_execute(op, index, task, outcome) return - if ( - self._context.reset_deadline is not None - and time.monotonic() >= self._context.reset_deadline - ): - self._finish_reset(False, "fault reset deadline exceeded") + if method == ActionMethod.CANCEL: + self._done_cancel(op, index, task, outcome) return - if any(task.action is not None for task in op.tasks): + if method == ActionMethod.STATUS: + self._done_status(op, index, task, outcome) return + def _aux_done(self, done: AuxiliaryDone) -> None: + self._auxiliary.complete(done) + self._condition.notify_all() + self._shutdown_progress() + + def _reset_next(self, op: Operation) -> None: + reset = self._context.reset_handle + if reset is None or reset in dict(self._context.reset_results): + return + if self._gate == "sealed": + self._finish_reset(False, "reset unavailable while runtime is closing") + self._shutdown_progress() + return proven = self._context.reset_proven_inactive - if event is not None and event.method == ActionMethod.STATUS: - if event.outcome in (Outcome.INACTIVE, Outcome.CANCELLED, Outcome.COMPLETED): - proven = proven | {event.task_id} if event.task_id is not None else proven - elif event.outcome in (Outcome.RUNNING, Outcome.ACCEPTED): - task = next(task for task in op.tasks if task.task_id == event.task_id) - self._schedule_action(self._context.active or op, task, ActionMethod.CANCEL) - return - else: - self._finish_reset(False, "status reconciliation is unsafe") - return - elif event is not None and event.method == ActionMethod.RESET: - if event.outcome != Outcome.INACTIVE: - self._finish_reset(False, "remote reset failed") - return - if event.task_id is not None: - self._commit( - replace( - self._context, - reset_completed_tasks=self._context.reset_completed_tasks | {event.task_id}, - ) - ) - elif event is not None and event.outcome in (Outcome.UNKNOWN, Outcome.FAILED): - self._finish_reset(False, "fault reconciliation is unsafe") + completed = self._context.reset_completed_tasks + for task in op.tasks: + if task.action is not None or task.activity == TaskActivity.NOT_STARTED: + continue + if task.task_id in completed: + continue + if task.task_id in proven: + if task.reset_required: + self._start_action(op, task, ActionMethod.RESET) + return + self._commit(reset_completed_tasks=frozenset((*completed, task.task_id))) + completed = self._context.reset_completed_tasks + continue + self._start_action(op, task, ActionMethod.STATUS) return + if all(t.task_id in completed or t.activity == TaskActivity.NOT_STARTED for t in op.tasks): + self._finish_reset(True) - if proven != self._context.reset_proven_inactive: - self._commit(replace(self._context, reset_proven_inactive=proven)) - - current = self._context.active or op - unproven = next( - ( - task - for task in current.tasks - if task.task_id not in self._context.reset_proven_inactive - ), - None, + def _done_reset( + self, + op: Operation, + index: int, + task: TaskRecord, + method: ActionMethod, + outcome: Outcome, + ) -> None: + current = op.tasks[index] + decision = reconcile_reset_completion( + method=method, + outcome=outcome, + activity=current.activity, + gate_sealed=self._gate == "sealed", + clock_failed=self._clock_invalid, ) - if unproven is not None: - method = ( - ActionMethod.CANCEL - if unproven.activity - in ( - TaskActivity.NOT_STARTED, - TaskActivity.EXECUTE_UNRESOLVED, - TaskActivity.ACTIVE, - TaskActivity.UNKNOWN, - ) - else ActionMethod.STATUS - ) - self._schedule_action(current, unproven, method) - return - - pending_reset = next( - ( - task - for task in current.tasks - if task.reset_required and task.task_id not in self._context.reset_completed_tasks + current = replace( + current, + activity=decision.activity or current.activity, + cancel_required=( + current.cancel_required + if decision.cancel_required is None + else decision.cancel_required + ), + reset_required=( + current.reset_required + if decision.reset_required is None + else decision.reset_required ), - None, ) - if pending_reset is not None: - self._schedule_action(current, pending_reset, ActionMethod.RESET) - return - self._finish_reset(True) - - def _reduce(self, event: _Event) -> Any: - """Pure transition function plus immutable context commit.""" - c = self._context - if c.shutdown == ShutdownState.CLOSED: - return False - if c.shutdown == ShutdownState.CLOSING and event.kind not in ( - "action_deadline", - "action_finished", - "effect_rejected", - "shutdown", - "shutdown_deadline", - ): - return False - if event.kind == "start_planning" and c.state in ( - LifecycleState.IDLE, - LifecycleState.READY, - ): - token = str(uuid.uuid4()) + self._commit(active=replace(op, tasks=(*op.tasks[:index], current, *op.tasks[index + 1 :]))) + if decision.prove_inactive: self._commit( - replace( - c, - state=LifecycleState.PLANNING, - planning_token=token, - ready_plan=None, - ready_plan_id=None, - diagnostic=None, + reset_proven_inactive=frozenset( + (*self._context.reset_proven_inactive, current.task_id) ) ) - return token - if event.kind == "cancel_planning" and c.state == LifecycleState.PLANNING: - self._commit(replace(c, state=LifecycleState.IDLE, planning_token=None)) - return True - if ( - event.kind == "plan_ready" - and c.state == LifecycleState.PLANNING - and event.payload is not None - ): - plan_id = str(uuid.uuid4()) + if decision.complete_reset: self._commit( - replace( - c, - state=LifecycleState.READY, - ready_plan=event.payload, - ready_plan_id=plan_id, - planning_token=None, - diagnostic=None, + reset_completed_tasks=frozenset( + (*self._context.reset_completed_tasks, current.task_id) ) ) - return plan_id - if event.kind == "planning_failed" and c.state == LifecycleState.PLANNING: - self._commit( - replace( - c, - state=LifecycleState.IDLE, - planning_token=None, - ready_plan=None, - ready_plan_id=None, - diagnostic=event.diagnostic or "planning failed", - ) + diagnostic = decision.diagnostic or ( + f"task={current.task_name} action={method.value} outcome={outcome.value}" + ) + if decision.reset_success is not None: + self._finish_reset(decision.reset_success, diagnostic) + if decision.request_cancel: + self._schedule_cancels( + self._context.active or op, + emergency=decision.emergency_cancel, ) - return True - if event.kind == "clear_ready_plan" and c.state == LifecycleState.READY: - self._commit(replace(c, state=LifecycleState.IDLE, ready_plan=None, ready_plan_id=None)) - return True - if event.kind == "dispatch_start" and c.state in ( - LifecycleState.IDLE, - LifecycleState.READY, - ): - handle, plan = event.payload + if decision.advance_reset: + self._reset_next(self._context.active or op) + self._shutdown_progress() + + def _record_dispatch(self, op: Operation, outcome: Outcome, diagnostic: str = "") -> None: + if dict(self._context.dispatch_results).get(op.handle) is None: self._commit( - replace( - c, - state=LifecycleState.DISPATCHING, - ready_plan=None, - ready_plan_id=None, - active=Operation( - handle, - plan, - tuple( - TaskRecord(str(uuid.uuid4()), e.task_name, e, TaskActivity.NOT_STARTED) - for e in plan.entries - ), - ), - fault=None, - diagnostic=None, - reset_handle=None, - reset_proven_inactive=frozenset(), - reset_completed_tasks=frozenset(), - reset_deadline=None, - shutdown_result=None, - shutdown_deadline=None, - ) + dispatch_results=( + *self._context.dispatch_results, + (op.handle, ExecutionResult(op.handle, outcome, diagnostic)), + )[-16:] ) - return True - if event.kind == "reset_deadline": - reset_handle = self._context.reset_handle - if ( - self._reset_pending() - and reset_handle is not None - and event.payload == reset_handle.reset_id - ): - self._finish_reset(False, "fault reset deadline exceeded") - return True - return False - if event.kind == "shutdown_deadline": - if ( - c.shutdown == ShutdownState.CLOSING - and c.shutdown_result is None - and event.payload == c.shutdown_deadline - ): - self._finish_shutdown(False, "shutdown safety deadline exceeded") - return True - return False - op = c.active - if event.kind == "effect_rejected": - if ( - op is None - or event.operation_id != op.handle.operation_id - or event.attempt_id != op.handle.attempt_id - ): - return False - task_index = next( - (i for i, task in enumerate(op.tasks) if task.task_id == event.task_id), None + + def _begin_cancellation(self, diagnostic: str = "") -> None: + op = self._context.active + if op is None: + return + tasks = tuple( + replace( + t, + cancel_required=( + t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), ) - if task_index is None: - return False - task = op.tasks[task_index] - if ( - task.action is None - or task.action.action_id != event.action_id - or task.action.method != event.method - ): - return False - if event.reset_id is not None: - current_reset_id = c.reset_handle.reset_id if c.reset_handle is not None else None - if event.reset_id != current_reset_id: - tasks = list(op.tasks) - tasks[task_index] = replace(task, action=None) - self._commit(replace(c, active=replace(op, tasks=tuple(tasks)))) - if self._reset_pending(): - self._reconcile_reset(self._context.active) - return True - activity = ( - TaskActivity.INACTIVE if event.method == ActionMethod.EXECUTE else task.activity + for t in op.tasks + ) + op = replace( + op, + tasks=tasks, + cancel_requested=True, + failed=op.failed or diagnostic == "physical operation deadline exceeded", + diagnostic=diagnostic or op.diagnostic, + ) + self._commit(state=LifecycleState.CANCELLING, active=op) + self._schedule_cancels(op) + + def _done_execute(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + if self._context.shutdown == ShutdownState.CLOSING: + op = replace(op, cancel_requested=True) + decision = reconcile_execute_completion(outcome) + if decision.fault: + diagnostic = f"task={task.task_name} action=execute outcome={outcome.value}" + self._record_dispatch(op, Outcome.UNKNOWN, diagnostic) + self._fault( + diagnostic, + op.handle, + replace(task, action=None), ) - tasks = list(op.tasks) - tasks[task_index] = replace(task, action=None, activity=activity) - self._commit(replace(c, active=replace(op, tasks=tuple(tasks)))) - if c.shutdown == ShutdownState.CLOSING and c.shutdown_result is None: - self._finish_shutdown(False, "shutdown dropped an unresolved coordinator effect") - return True - if event.kind in ("action_deadline", "action_finished"): - if ( - op is None - or event.operation_id != op.handle.operation_id - or event.attempt_id != op.handle.attempt_id - or event.task_id not in {t.task_id for t in op.tasks} - ): - return False - task_index = next(i for i, task in enumerate(op.tasks) if task.task_id == event.task_id) - task = op.tasks[task_index] - if ( - task.action is None - or task.action.action_id != event.action_id - or task.action.method != event.method - ): - return False - if event.reset_id is not None: - current_reset_id = c.reset_handle.reset_id if c.reset_handle is not None else None - if event.reset_id != current_reset_id: - if event.kind == "action_deadline": - return False - tasks = list(op.tasks) - tasks[task_index] = replace(task, action=None) - refreshed = replace(op, tasks=tuple(tasks)) - self._commit(replace(c, active=refreshed)) - self._reconcile_reset(refreshed) - return True - event_detail = f"task={task.task_name} action={event.action_id} outcome={event.outcome.value if event.outcome is not None else 'unknown'}" - if event.kind == "action_deadline": - if task.action.deadline_reported: - return False - if not self._action_was_started(task.action.action_id): - return self._reduce(replace(event, kind="effect_rejected")) - action = replace(task.action, deadline_reported=True) - tasks = list(op.tasks) - tasks[task_index] = replace( - task, action=action, activity=TaskActivity.UNKNOWN, cancel_required=True - ) - with self._gateway_condition: - self._started_actions.discard(task.action.action_id) - fault_op = replace(op, tasks=tuple(tasks), uncertain=True) - result = ExecutionResult(op.handle, Outcome.UNKNOWN, "coordinator action deadline") - self._commit( - replace( - c, - state=LifecycleState.FAULT, - fault="coordinator action deadline", - diagnostic="coordinator action deadline", - active=fault_op, - ) - ) - if c.shutdown == ShutdownState.CLOSING: - self._finish_shutdown(False, "shutdown safety unresolved: action deadline") - return True - if self._reset_pending(): - self._finish_reset(False, "fault reconciliation action deadline") - return True - self._record_dispatch_result_if_absent(result) - self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) - return True - tasks = list(op.tasks) - with self._gateway_condition: - self._started_actions.discard(task.action.action_id) - if event.outcome in (Outcome.ACCEPTED, Outcome.RUNNING): - activity = TaskActivity.ACTIVE - elif event.outcome == Outcome.COMPLETED: - activity = TaskActivity.COMPLETED - elif event.outcome == Outcome.CANCELLED: - activity = ( - TaskActivity.CANCELLED - if event.method == ActionMethod.CANCEL or c.state == LifecycleState.CANCELLING - else TaskActivity.UNKNOWN + return + if decision.rejected: + tasks = tuple( + replace( + t, + action=None, + activity=TaskActivity.INACTIVE if i >= index else t.activity, + cancel_required=(t.cancel_required or i < index) if i < index else False, ) - elif event.outcome == Outcome.INACTIVE: - activity = TaskActivity.INACTIVE - elif event.outcome == Outcome.FAILED: - activity = TaskActivity.REMOTE_FAULT - elif event.outcome == Outcome.REJECTED and event.method == ActionMethod.EXECUTE: - activity = TaskActivity.INACTIVE + for i, t in enumerate(op.tasks) + ) + op = replace(op, tasks=tasks, rejected=True, cancel_requested=True) + diagnostic = f"task={task.task_name} action=execute outcome={outcome.value}" + op = replace(op, diagnostic=diagnostic) + if self._context.state == LifecycleState.FAULT or op.uncertain: + self._commit(active=op) else: - activity = TaskActivity.UNKNOWN - tasks[task_index] = replace( - task, - activity=activity, + self._commit(state=LifecycleState.CANCELLING, active=op) + self._record_dispatch(op, Outcome.REJECTED, diagnostic) + self._schedule_cancels(op) + self._finish_if_terminal(op) + return + activity = TaskActivity.ACTIVE + tasks = tuple( + replace( + t, action=None, - cancel_required=task.cancel_required - or (c.state in (LifecycleState.CANCELLING, LifecycleState.FAULT)), - reset_required=task.reset_required or event.outcome == Outcome.FAILED, + activity=activity, + cancel_required=( + t.cancel_required + or op.cancel_requested + or self._context.state == LifecycleState.FAULT + ), ) - next_op = replace(op, tasks=tuple(tasks)) - self._commit(replace(c, active=next_op)) - if c.shutdown == ShutdownState.CLOSING: - if self._context.shutdown_result is None: - self._maybe_finish_shutdown() - return True - if self._reset_pending(): - self._reconcile_reset(next_op, event) - return True - if c.state == LifecycleState.FAULT: - if event.method != ActionMethod.CANCEL: - self._schedule_actions(next_op, next_op.tasks, ActionMethod.CANCEL) - return True - if event.method == ActionMethod.STATUS and event.outcome in ( - Outcome.UNKNOWN, - Outcome.FAILED, - ): - fault_op = replace(next_op, uncertain=True) - diagnostic = f"status is unsafe: {event_detail}" - result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) - self._commit( - replace( - self._context, - state=LifecycleState.FAULT, - fault=diagnostic, - diagnostic=diagnostic, - active=fault_op, - ) - ) - self._record_dispatch_result_if_absent(result) - self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) - return True - if event.method == ActionMethod.CANCEL and event.outcome in ( - Outcome.UNKNOWN, - Outcome.FAILED, - ): - fault_op = replace(next_op, uncertain=True) - diagnostic = ( - "cancellation is uncertain" - if event.outcome == Outcome.UNKNOWN - else "cancellation failed" - ) - result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) - self._commit( - replace( - self._context, - state=LifecycleState.FAULT, - fault=diagnostic, - diagnostic=diagnostic, - active=fault_op, - ) - ) - self._record_dispatch_result_if_absent(result) - if event.outcome == Outcome.FAILED: - self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) - return True - if ( - c.state == LifecycleState.CANCELLING - and event.method == ActionMethod.STATUS - and event.outcome == Outcome.RUNNING - ): - self._schedule_actions(next_op, (tasks[task_index],), ActionMethod.CANCEL) - return True - if ( - event.method == ActionMethod.STATUS - and event.outcome in (Outcome.INACTIVE, Outcome.CANCELLED) - and c.state == LifecycleState.RUNNING - ): - diagnostic = f"status reported terminal failure: {event_detail}" - failed_op = replace( - next_op, - cancel_requested=True, - cleanup_required=True, - rejected=False, - failed=True, - diagnostic=diagnostic, - ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + self._commit(active=op) + if op.cancel_requested or self._context.state in ( + LifecycleState.CANCELLING, + LifecycleState.FAULT, + ): + self._record_dispatch(op, Outcome.CANCELLED) + self._schedule_cancels(op) + self._shutdown_progress() + return + if index != op.next_index: + self._fault("out-of-order execute completion", op.handle, task) + return + if index + 1 < len(op.tasks): + op = replace(op, next_index=index + 1) + self._commit(active=op) + self._start_action(op, op.tasks[index + 1], ActionMethod.EXECUTE) + else: + self._record_dispatch(op, Outcome.ACCEPTED) + now = self._clock_now() + if now is None: + return + self._commit( + state=LifecycleState.RUNNING, + physical_operation_deadline=now + self._physical_timeout, + ) + self._next_poll = now + self._poll_interval + + def _done_cancel(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + decision = reconcile_cancel_completion(outcome) + if decision.fault and outcome in (Outcome.UNKNOWN, Outcome.FAILED): + diagnostic = decision.diagnostic or "cancellation failed" + self._fault( + diagnostic, + op.handle, + replace(task, action=None), + schedule_cancels=False, + mark_peers=False, + ) + current = self._context.active + if current is not None and current.handle == op.handle: self._commit( - replace( - self._context, - state=LifecycleState.CANCELLING, - active=failed_op, - diagnostic=diagnostic, + active=replace( + current, + tasks=tuple( + replace( + item, + action=None, + cancel_required=False, + reset_required=True, + ) + if item.task_id == task.task_id + else item + for item in current.tasks + ), ) ) - self._schedule_actions(failed_op, failed_op.tasks, ActionMethod.CANCEL) - return True - if event.method == ActionMethod.EXECUTE and event.outcome in ( - Outcome.REJECTED, - Outcome.UNKNOWN, - Outcome.FAILED, - ): - if event.outcome == Outcome.REJECTED: - skipped = tuple( - replace(item, activity=TaskActivity.INACTIVE) - if item.activity == TaskActivity.NOT_STARTED + current = self._context.active + if current is not None: + if self._context.shutdown == ShutdownState.CLOSING: + self._finish_shutdown(False, diagnostic) + self._schedule_cancels(current) + self._shutdown_progress() + return + if decision.fault: + diagnostic = f"malformed cancel outcome for {task.task_name}" + self._fault( + diagnostic, + op.handle, + replace(task, action=None), + schedule_cancels=False, + mark_peers=False, + ) + current = self._context.active + if current is not None and current.handle == op.handle: + current = replace( + current, + tasks=tuple( + replace(item, action=None, cancel_required=False) + if item.task_id == task.task_id else item - for item in next_op.tasks - ) - next_op = replace(next_op, tasks=skipped) - self._commit(replace(self._context, active=next_op)) - if event.outcome != Outcome.REJECTED: - fault_op = replace(next_op, uncertain=True) - diagnostic = f"execute outcome is uncertain: {event_detail}" - result = ExecutionResult(op.handle, Outcome.UNKNOWN, diagnostic) - self._commit( - replace( - self._context, - state=LifecycleState.FAULT, - fault=diagnostic, - diagnostic=diagnostic, - active=fault_op, - ) - ) - self._record_dispatch_result_if_absent(result) - self._schedule_actions(fault_op, fault_op.tasks, ActionMethod.CANCEL) - elif any( - t.activity in (TaskActivity.ACTIVE, TaskActivity.UNKNOWN) for t in next_op.tasks - ): - diagnostic = f"execute rejected: {event_detail}" - rejected_op = replace( - next_op, - cancel_requested=True, - rejected=True, - cleanup_required=True, - diagnostic=diagnostic, - ) - rejection_result = ExecutionResult(op.handle, Outcome.REJECTED, diagnostic) - self._commit( - replace( - self._context, - state=LifecycleState.CANCELLING, - active=rejected_op, - diagnostic=diagnostic, - ) - ) - self._record_dispatch_result_if_absent(rejection_result) - self._schedule_actions(rejected_op, rejected_op.tasks, ActionMethod.CANCEL) - else: - diagnostic = f"execute rejected: {event_detail}" - result = ExecutionResult(next_op.handle, Outcome.REJECTED, diagnostic) - self._commit( - replace( - self._context, - state=LifecycleState.IDLE, - active=None, - diagnostic=diagnostic, - terminal_results=( - *self._context.terminal_results, - (next_op.handle, result), - )[-16:], - ) - ) - self._record_dispatch_result_if_absent(result) - return True - self._advance(next_op, event) - return True - if ( - ( - event.kind == "cancel" - or ( - event.kind == "cancel_if_current" - and op is not None - and event.payload == op.handle + for item in current.tasks + ), ) + self._commit(active=current) + if self._context.shutdown == ShutdownState.CLOSING: + self._finish_shutdown(False, diagnostic) + self._schedule_cancels(current) + self._shutdown_progress() + return + tasks = tuple( + replace(t, action=None, activity=TaskActivity.CANCELLED, cancel_required=False) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + self._commit(active=op) + self._schedule_cancels(op) + self._finish_if_terminal(op) + self._shutdown_progress() + + def _done_status(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + decision = reconcile_status_completion(outcome) + if decision.fault: + self._fault( + f"status is unsafe: task={task.task_name} action=get_state outcome={outcome.value}", + op.handle, + replace(task, action=None), ) - and op is not None - and c.state in (LifecycleState.DISPATCHING, LifecycleState.RUNNING) - ): - cancel_tasks: tuple[TaskRecord, ...] = tuple( - replace(task, activity=TaskActivity.INACTIVE, cancel_required=True) - if task.activity == TaskActivity.NOT_STARTED - else replace(task, cancel_required=True) - if task.activity == TaskActivity.EXECUTE_UNRESOLVED - else task - for task in op.tasks + return + activity = cast("TaskActivity", decision.activity) + tasks = tuple( + replace( + t, + action=None, + activity=activity, + cancel_required=False + if t.task_id == task.task_id + and activity + in (TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED) + else t.cancel_required, ) - next_op = replace(op, tasks=cancel_tasks, cancel_requested=True) - self._commit(replace(c, state=LifecycleState.CANCELLING, active=next_op)) - if c.state == LifecycleState.DISPATCHING and not dict(c.dispatch_results).get( - op.handle - ): - result = ExecutionResult( - op.handle, Outcome.CANCELLED, "cancelled before dispatch completed" - ) - self._record_dispatch_result_if_absent(result) - self._schedule_actions(next_op, cancel_tasks, ActionMethod.CANCEL) - return True - if event.kind == "poll_tick" and op is not None and c.state == LifecycleState.RUNNING: - candidate = next( - ( - item - for item in op.tasks - if item.activity == TaskActivity.ACTIVE and item.action is None + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + if activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED): + op = replace( + op, + cancel_requested=True, + failed=True, + diagnostic=( + f"task={task.task_name} outcome={outcome.value}" + if activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED) + else op.diagnostic + ), + tasks=tuple( + replace( + t, + cancel_required=( + False + if t.task_id == task.task_id + else t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + ) + for t in op.tasks ), + ) + self._commit(state=LifecycleState.CANCELLING, active=op) + self._schedule_cancels(op) + self._finish_if_terminal(op) + return + self._commit(active=op) + self._finish_if_terminal(op) + if self._context.state in (LifecycleState.CANCELLING, LifecycleState.FAULT): + self._schedule_cancels(op) + + def _poll_active(self) -> None: + op = self._context.active + if op is not None and self._context.state == LifecycleState.RUNNING: + task = next( + (t for t in op.tasks if t.activity == TaskActivity.ACTIVE and t.action is None), None, ) - if candidate is not None: - self._schedule_actions(op, (candidate,), ActionMethod.STATUS) - return True - if event.kind == "fault": - diagnostic = event.diagnostic or "remote uncertainty" - self._commit( - replace( - c, - state=LifecycleState.FAULT, - fault=diagnostic, - diagnostic=diagnostic, - active=replace(op, uncertain=True) if op else None, + if task: + self._start_action(op, task, ActionMethod.STATUS) + + def _schedule_cancels(self, op: Operation, *, emergency: bool = False) -> None: + for task in op.tasks: + current = self._context.active + if current is None or current.handle != op.handle: + return + task = next(t for t in current.tasks if t.task_id == task.task_id) + if task.action is None and task.cancel_required: + self._start_action( + current, task, ActionMethod.CANCEL, emergency=emergency or self._clock_invalid ) - ) - return True - if event.kind == "shutdown": - if self._reset_pending(): - self._finish_reset(False, "shutdown interrupted fault reset") - c = self._context - op = c.active - if op is not None: - shutdown_tasks: tuple[TaskRecord, ...] = tuple( - replace(task, activity=TaskActivity.INACTIVE, cancel_required=True) - if task.activity == TaskActivity.NOT_STARTED - else replace(task, cancel_required=True) - for task in op.tasks + + def _finish_if_terminal(self, op: Operation) -> None: + if any(t.action is not None for t in op.tasks): + return + if self._context.state == LifecycleState.CANCELLING: + if all( + t.activity + in ( + TaskActivity.NOT_STARTED, + TaskActivity.CANCELLED, + TaskActivity.INACTIVE, + TaskActivity.COMPLETED, ) - op = replace(op, tasks=shutdown_tasks, cancel_requested=True, cleanup_required=True) - self._commit( + and not t.cancel_required + for t in op.tasks + ): + if self._context.shutdown == ShutdownState.CLOSING: + self._shutdown_progress() + else: + self._terminal( + Outcome.REJECTED + if op.rejected and not op.failed + else Outcome.FAILED + if op.failed + else Outcome.CANCELLED + ) + elif all(t.activity == TaskActivity.COMPLETED for t in op.tasks): + self._terminal(Outcome.COMPLETED) + + def _terminal(self, outcome: Outcome) -> None: + op = self._context.active + if op is None: + return + result = ExecutionResult(op.handle, outcome, op.diagnostic) + self._commit( + state=LifecycleState.IDLE, + active=None, + diagnostic=op.diagnostic, + terminal_results=(*self._context.terminal_results, (op.handle, result))[-16:], + ) + + def _fault( + self, + diagnostic: str, + handle: OperationHandle | None = None, + task: TaskRecord | None = None, + *, + schedule_cancels: bool = True, + mark_peers: bool = True, + ) -> None: + op = self._context.active + if op is not None: + tasks = tuple( replace( - c, - shutdown=ShutdownState.CLOSING, - state=LifecycleState.CANCELLING if op else LifecycleState.IDLE, - planning_token=None, - ready_plan=None, - ready_plan_id=None, - active=op, - shutdown_deadline=event.payload, + t, + action=None if task is not None and task.action is None else t.action, + activity=TaskActivity.UNKNOWN, + cancel_required=( + t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + reset_required=( + t.reset_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), ) - ) - if op is None: - self._finish_shutdown(True) - else: - self._schedule_actions(op, op.tasks, ActionMethod.CANCEL) - self._maybe_finish_shutdown() - return True - return False - - def _advance(self, op: Operation, event: _Event) -> None: - # This is called only by the owner after a correlated result. - cancelling = op.cancel_requested or self._context.state in ( - LifecycleState.CANCELLING, - LifecycleState.FAULT, - ) - if cancelling: - self._schedule_actions(op, op.tasks, ActionMethod.CANCEL) - if ( - not cancelling - and event.method == ActionMethod.EXECUTE - and event.outcome == Outcome.ACCEPTED - ): - if op.next_index + 1 < len(op.tasks): - self._commit( - replace(self._context, active=replace(op, next_index=op.next_index + 1)) + if ( + mark_peers + and t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) ) - self._schedule_execute(self._context.active) - return - if all(task.activity == TaskActivity.ACTIVE for task in op.tasks): - result = ExecutionResult(op.handle, Outcome.ACCEPTED) - self._commit(replace(self._context, state=LifecycleState.RUNNING)) - self._record_dispatch_result_if_absent(result) - return - terminal = all( - task.action is None - and task.activity - in (TaskActivity.COMPLETED, TaskActivity.CANCELLED, TaskActivity.INACTIVE) - for task in op.tasks - ) - completed = all( - task.activity == TaskActivity.COMPLETED and task.action is None for task in op.tasks - ) - if terminal and (op.cancel_requested or completed): - if self._context.state == LifecycleState.FAULT: - return - outcome = ( - Outcome.REJECTED - if op.rejected - else Outcome.FAILED - if op.failed - else Outcome.CANCELLED - if op.cancel_requested - else Outcome.COMPLETED + or (task is not None and t.task_id == task.task_id) + else t + for t in op.tasks ) - result = ExecutionResult(op.handle, outcome, op.diagnostic) + op = replace(op, tasks=tasks, uncertain=True, diagnostic=diagnostic) + self._commit( + state=LifecycleState.FAULT, + fault=self._context.fault or diagnostic, + diagnostic=diagnostic, + active=op, + ) + fault_handle = handle or (op.handle if op is not None else None) + if fault_handle is not None: + self._fault_terminal_results[fault_handle] = diagnostic + if len(self._fault_terminal_results) > 8: + oldest = next(iter(self._fault_terminal_results)) + del self._fault_terminal_results[oldest] + if handle and not dict(self._context.dispatch_results).get(handle): self._commit( - replace( - self._context, - state=LifecycleState.IDLE, - active=None, - diagnostic=op.diagnostic or self._context.diagnostic, - terminal_results=(*self._context.terminal_results, (op.handle, result))[-16:], - ) + dispatch_results=( + *self._context.dispatch_results, + (handle, ExecutionResult(handle, Outcome.UNKNOWN, diagnostic)), + )[-16:] ) + if op and schedule_cancels: + self._schedule_cancels(op) - def _schedule_execute(self, op: Operation | None) -> None: - if op is None or op.next_index >= len(op.tasks): + def _stop_done(self, done: StopDone) -> None: + self._shutdown_stop_finished = True + if self._context.shutdown_result is not None: + self._shutdown_drain_owner() return - task = op.tasks[op.next_index] - if task.action is not None: + if not done.success: + self._finish_shutdown( + False, + f"gateway close failed: {done.diagnostic or 'stop failed'}", + ) + self._shutdown_drain_owner() return - self._schedule_action(op, task, ActionMethod.EXECUTE) - - def _schedule_action(self, op: Operation, task: TaskRecord, method: ActionMethod) -> None: - self._schedule_actions(op, (task,), method) - - def _schedule_actions(self, op: Operation, requested: Any, method: ActionMethod) -> None: - """Allocate and commit a whole cleanup batch against one operation.""" - if ( - self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED) - or self._context.shutdown == ShutdownState.CLOSED - or (self._context.shutdown == ShutdownState.CLOSING and method != ActionMethod.CANCEL) - ): + deadline = self._context.shutdown_deadline + now = self._clock_now() + if now is None: return - candidates = tuple( - task - for task in requested - if task.action is None - and ( - method in (ActionMethod.STATUS, ActionMethod.RESET) - or task.activity - not in ( - TaskActivity.INACTIVE, - TaskActivity.COMPLETED, - TaskActivity.CANCELLED, - TaskActivity.REMOTE_FAULT, - ) - ) - ) - if not candidates: + if deadline is None or now >= deadline: + self._finish_shutdown(False, "shutdown safety deadline exceeded") + self._shutdown_drain_owner() return - now = time.monotonic() - reset_id = ( - self._context.reset_handle.reset_id - if self._reset_pending() and self._context.reset_handle is not None - else None + self._commit( + state=LifecycleState.IDLE, + shutdown=ShutdownState.CLOSED, + shutdown_result=ShutdownResult(True, done.diagnostic), + shutdown_deadline=None, ) - actions = { - task.task_id: ActionRecord( - str(uuid.uuid4()), method, now, now + self._timeout, reset_id=reset_id + self._shutdown_drain_owner() + + def _has_unsettled_shutdown_work(self) -> bool: + active = self._context.active + return bool( + self._pending + or self._auxiliary.has_unsettled() + or ( + active is not None + and any(t.action is not None or t.cancel_required for t in active.tasks) ) - for task in candidates - } - task_map = {task.task_id: task for task in op.tasks} - for task in candidates: - task_map[task.task_id] = replace( - task, - action=actions[task.task_id], - activity=TaskActivity.EXECUTE_UNRESOLVED - if method == ActionMethod.EXECUTE - else task.activity, - ) - updated = replace(op, tasks=tuple(task_map[task.task_id] for task in op.tasks)) - self._commit(replace(self._context, active=updated)) - self._effects.extend((updated, task_map[task.task_id], method) for task in candidates) - - def _process_effects(self) -> None: - effects, self._effects = self._effects, [] - if self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED): - if ( - self._context.shutdown == ShutdownState.CLOSING - and self._context.shutdown_result is None - and effects - ): - self._finish_shutdown(False, "shutdown gateway closed before pending effects") + ) + + def _shutdown_progress(self) -> None: + if self._context.shutdown != ShutdownState.CLOSING: return - for op, task, method in effects: - admitted = self._run_action(op, task, method) - if ( - not admitted - and self._context.shutdown == ShutdownState.CLOSING - and self._context.shutdown_result is None - ): - self._finish_shutdown(False, "shutdown dropped an unresolved coordinator effect") - - def _invoke_gateway_action( - self, action_id: str, task_name: str, method: ActionMethod, request: Any - ) -> tuple[bool, Any]: - """Lease a gateway call immediately before invoking it; no context lock crosses I/O.""" - with self._gateway_condition: - allowed = self._gateway_gate == _GatewayGateState.OPEN or ( - self._gateway_gate == _GatewayGateState.CANCEL_ONLY - and method == ActionMethod.CANCEL - ) - if not allowed: - return False, None - self._rpc_inflight += 1 - self._admitted_actions.add(action_id) - self._started_actions.add(action_id) - try: - outcome: Any = Outcome.UNKNOWN - try: - if method == ActionMethod.EXECUTE: - outcome = self._gateway.execute(task_name, request) - elif method == ActionMethod.CANCEL: - outcome = self._gateway.cancel(task_name) - elif method == ActionMethod.STATUS: - outcome = self._gateway.status(task_name) - elif method == ActionMethod.RESET: - outcome = self._gateway.reset(task_name) - elif method == ActionMethod.GRIPPER_SET: - outcome = self._gateway.set_gripper_position(task_name, float(request)) - else: - outcome = self._gateway.get_gripper_position(task_name) - except Exception: - outcome = Outcome.UNKNOWN - return True, outcome - finally: - with self._gateway_condition: - self._rpc_inflight -= 1 - self._admitted_actions.discard(action_id) - self._gateway_condition.notify_all() - - def _action_was_admitted(self, action_id: str) -> bool: - with self._gateway_condition: - return action_id in self._admitted_actions - - def _action_was_started(self, action_id: str) -> bool: - with self._gateway_condition: - return action_id in self._started_actions - - def _run_action(self, op: Operation, task: TaskRecord, method: ActionMethod) -> bool: - action = task.action - if action is None: - return False - request = task.entry.request + if self._has_unsettled_shutdown_work(): + return + if self._shutdown_stop_started: + return + self._shutdown_stop_started = True + self._gate = "sealed" + self._effects.submit_stop(self._gateway.stop, self._events.put) - def call(_cancel: threading.Event) -> None: - admitted, outcome = self._invoke_gateway_action( - action.action_id, task.task_name, method, request - ) - if admitted: - self._events.put( - _Call( - lambda: self._reduce( - _Event( - "action_finished", - op.handle.operation_id, - op.handle.attempt_id, - task.task_id, - action.action_id, - method, - outcome, - reset_id=action.reset_id, - ) - ) - ) - ) - else: - self._events.put( - _Call( - lambda: self._reduce( - _Event( - "effect_rejected", - op.handle.operation_id, - op.handle.attempt_id, - task.task_id, - action.action_id, - method, - reset_id=action.reset_id, - ) - ) - ) - ) + def _shutdown_drain_owner(self) -> None: + if self._shutdown_drain_started: + return + if self._has_unsettled_shutdown_work(): + return + self._shutdown_drain_started = True + self._effects.shutdown() + self._events.put(None) - def deadline(cancel: threading.Event) -> None: - if cancel.wait(max(0.0, action.deadline - time.monotonic())): - return - self._events.put( - _Call( - lambda: self._reduce( - _Event( - "action_deadline", - op.handle.operation_id, - op.handle.attempt_id, - task.task_id, - action.action_id, - method, - reset_id=action.reset_id, - ) - ) - ) - ) + def _finish_reset(self, success: bool, diagnostic: str = "") -> None: + h = self._context.reset_handle + if h is None or h in dict(self._context.reset_results): + return + result = ResetResult(h, success, diagnostic) + self._commit( + state=LifecycleState.IDLE if success else LifecycleState.FAULT, + reset_results=(*self._context.reset_results, (h, result))[-16:], + reset_handle=None if success else h, + reset_deadline=None, + fault=None if success else self._context.fault or diagnostic, + active=None if success else self._context.active, + physical_operation_deadline=None + if success + else self._context.physical_operation_deadline, + ready_plan=None if success else self._context.ready_plan, + ready_plan_id=None if success else self._context.ready_plan_id, + diagnostic=diagnostic, + ) - self._start_auxiliary(call, name=f"action-{action.action_id}") - self._start_auxiliary(deadline, name=f"deadline-{action.action_id}") - return True + def _finish_shutdown(self, success: bool, diagnostic: str = "") -> None: + if self._context.shutdown_result is None: + self._commit( + shutdown_result=ShutdownResult(success, diagnostic), + shutdown_deadline=None if not success else self._context.shutdown_deadline, + ) def start_planning(self) -> str | None: - return cast("str | None", self._submit(lambda: self._reduce(_Event("start_planning")))) + def fn() -> str | None: + if self._context.state not in (LifecycleState.IDLE, LifecycleState.READY): + return None + token = str(uuid.uuid4()) + self._commit( + state=LifecycleState.PLANNING, + planning_token=token, + ready_plan=None, + ready_plan_id=None, + diagnostic=None, + ) + return token + + return self._submit(fn) def complete_planning(self, token: str, plan: PlanInput) -> CommandResult[str]: def fn() -> CommandResult[str]: - if self._context.planning_token != token: + if token != self._context.planning_token: return CommandResult( - False, diagnostic="planning token mismatch", snapshot=self.snapshot() + False, diagnostic="planning token mismatch", snapshot=self._snapshot() + ) + if not self._valid(plan): + self._commit( + state=LifecycleState.IDLE, + planning_token=None, + diagnostic="invalid prepared plan", ) - if not self._validate_plan(plan): - self._reduce(_Event("planning_failed", diagnostic="invalid prepared plan")) return CommandResult( - False, diagnostic="invalid prepared plan", snapshot=self.snapshot() + False, diagnostic="invalid prepared plan", snapshot=self._snapshot() ) - plan_id = self._reduce(_Event("plan_ready", payload=plan)) - return CommandResult(True, plan_id, snapshot=self.snapshot()) + pid = str(uuid.uuid4()) + self._commit( + state=LifecycleState.READY, ready_plan=plan, ready_plan_id=pid, planning_token=None + ) + return CommandResult(True, pid, snapshot=self._snapshot()) return self._submit(fn) def fail_planning(self, token: str, diagnostic: str) -> CommandResult[None]: - return self._submit( - lambda: CommandResult( - self._context.planning_token == token, - diagnostic=diagnostic - if self._context.planning_token == token - else "planning token mismatch", - snapshot=self.snapshot(), - ) - if self._context.planning_token != token - else ( - self._reduce(_Event("planning_failed", diagnostic=diagnostic)) - and CommandResult(True, diagnostic=diagnostic, snapshot=self.snapshot()) + def fn() -> CommandResult[None]: + ok = token == self._context.planning_token + if ok: + self._commit(state=LifecycleState.IDLE, planning_token=None, diagnostic=diagnostic) + return CommandResult( + ok, + diagnostic=diagnostic if ok else "planning token mismatch", + snapshot=self._snapshot(), ) - ) + + return self._submit(fn) def cancel_planning(self) -> CommandResult[None]: - return self._submit( - lambda: CommandResult( - bool(self._reduce(_Event("cancel_planning"))), snapshot=self.snapshot() - ) - ) + return self._submit(lambda: self._planning_command(LifecycleState.PLANNING, "cancelled")) + + def _planning_command(self, state: LifecycleState, diagnostic: str) -> CommandResult[None]: + ok = self._context.state == state + if ok: + self._commit(state=LifecycleState.IDLE, planning_token=None, diagnostic=diagnostic) + return CommandResult(ok, snapshot=self._snapshot()) def clear_ready_plan(self) -> CommandResult[None]: - return self._submit( - lambda: CommandResult( - bool(self._reduce(_Event("clear_ready_plan"))), snapshot=self.snapshot() + def fn() -> CommandResult[None]: + if self._context.state != LifecycleState.READY: + return CommandResult(False, snapshot=self._snapshot()) + self._commit( + state=LifecycleState.IDLE, + ready_plan=None, + ready_plan_id=None, + diagnostic="", ) - ) - - def execute_explicit(self, prepared_plan: PlanInput) -> CommandResult[OperationHandle]: - def fn() -> CommandResult[OperationHandle]: - if self._context.state in ( - LifecycleState.DISPATCHING, - LifecycleState.RUNNING, - LifecycleState.CANCELLING, - ) and self._validate_plan(prepared_plan): - handle = OperationHandle(str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) - result = ExecutionResult(handle, Outcome.REJECTED, "execution already active") - self._commit( - replace( - self._context, - diagnostic="execution already active", - terminal_results=(*self._context.terminal_results, (handle, result))[-16:], - ) - ) - return CommandResult( - False, - handle, - diagnostic="execution already active", - snapshot=self.snapshot(), - ) - if self._context.state not in ( - LifecycleState.IDLE, - LifecycleState.READY, - ) or not self._validate_plan(prepared_plan): - self._commit( - replace(self._context, diagnostic="execution unavailable or invalid plan") - ) - return CommandResult( - False, - diagnostic="execution unavailable or invalid plan", - snapshot=self.snapshot(), - ) - handle = OperationHandle(str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) - if self._reduce(_Event("dispatch_start", payload=(handle, prepared_plan))): - self._schedule_execute(self._context.active) - return CommandResult(True, handle, snapshot=self.snapshot()) - return CommandResult(False, diagnostic="dispatch rejected", snapshot=self.snapshot()) - - return self._submit(fn) - - def execute_ready(self) -> CommandResult[OperationHandle]: - def fn() -> CommandResult[OperationHandle]: - if self._context.ready_plan is None or self._context.ready_plan_id is None: - return CommandResult(False, diagnostic="no ready plan", snapshot=self.snapshot()) - plan = self._context.ready_plan - return self._execute_owner(plan, self._context.ready_plan_id) + return CommandResult(True, snapshot=self._snapshot()) return self._submit(fn) - def _execute_owner( + def _execute( self, plan: PlanInput, plan_id: str | None = None ) -> CommandResult[OperationHandle]: - handle = OperationHandle(plan_id or str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) - if self._reduce(_Event("dispatch_start", payload=(handle, plan))): - self._schedule_execute(self._context.active) - return CommandResult(True, handle, snapshot=self.snapshot()) - return CommandResult(False, diagnostic="dispatch rejected", snapshot=self.snapshot()) + if self._context.state not in ( + LifecycleState.IDLE, + LifecycleState.READY, + ) or not self._valid(plan): + return CommandResult( + False, diagnostic="execution unavailable or invalid plan", snapshot=self._snapshot() + ) + h = OperationHandle(plan_id or str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) + op = Operation( + h, + plan, + tuple( + TaskRecord(str(uuid.uuid4()), e.task_name, e, TaskActivity.NOT_STARTED) + for e in plan.entries + ), + ) + self._commit( + state=LifecycleState.DISPATCHING, + ready_plan=None, + ready_plan_id=None, + active=op, + diagnostic=None, + ) + self._start_action(op, op.tasks[0], ActionMethod.EXECUTE) + return CommandResult(True, h, snapshot=self._snapshot()) - def cancel(self) -> CommandResult[None]: + def execute_explicit(self, prepared_plan: PlanInput) -> CommandResult[OperationHandle]: + return self._submit(lambda: self._execute(prepared_plan)) + + def execute_ready(self) -> CommandResult[OperationHandle]: return self._submit( - lambda: CommandResult(bool(self._reduce(_Event("cancel"))), snapshot=self.snapshot()) + lambda: self._execute(self._context.ready_plan, self._context.ready_plan_id) + if self._context.ready_plan is not None + else CommandResult(False, diagnostic="no ready plan", snapshot=self._snapshot()) ) + def cancel(self) -> CommandResult[None]: + return self._submit(self._cancel_owner, closing=True) + + def _cancel_owner(self) -> CommandResult[None]: + op = self._context.active + if op is None or self._context.state not in ( + LifecycleState.DISPATCHING, + LifecycleState.RUNNING, + ): + return CommandResult(False, snapshot=self._snapshot()) + self._begin_cancellation() + return CommandResult(True, snapshot=self._snapshot()) + def cancel_if_current(self, handle: OperationHandle) -> CommandResult[None]: - """Cancel only if ``handle`` is the current operation at the owner event.""" + return self._submit( + lambda: self._cancel_owner() + if self._context.active and self._context.active.handle == handle + else CommandResult( + False, diagnostic="operation is no longer current", snapshot=self._snapshot() + ), + closing=True, + ) + def poll(self) -> CommandResult[None]: def fn() -> CommandResult[None]: - accepted = bool(self._reduce(_Event("cancel_if_current", payload=handle))) - current = self._context.active - diagnostic = ( - "" - if accepted - else "operation is no longer current" - if current is None or current.handle != handle - else f"cancellation rejected in state {self._context.state.name}" - ) + self._poll_active() return CommandResult( - accepted, - diagnostic=diagnostic, - snapshot=self.snapshot(), + self._context.state == LifecycleState.RUNNING, + snapshot=self._snapshot(), ) return self._submit(fn) - def poll(self) -> CommandResult[None]: - return self._submit( - lambda: CommandResult(bool(self._reduce(_Event("poll_tick"))), snapshot=self.snapshot()) - ) - def set_gripper_position(self, hardware_id: str, position: float) -> CommandResult[Outcome]: - def fn() -> CommandResult[Outcome]: - admitted, outcome = self._invoke_gateway_action( - str(uuid.uuid4()), hardware_id, ActionMethod.GRIPPER_SET, position - ) - if not admitted: - return CommandResult( - False, diagnostic="gateway is closing", snapshot=self.snapshot() - ) - return CommandResult(True, outcome, snapshot=self.snapshot()) - - return self._submit(fn) + result = self._submit(lambda: self._gripper(True, hardware_id, position)) + if isinstance(result, CommandResult): + return cast("CommandResult[Outcome]", result) + return cast("CommandResult[Outcome]", self._wait_aux(result)) def get_gripper_position(self, hardware_id: str) -> CommandResult[float]: - def fn() -> CommandResult[float]: - admitted, position = self._invoke_gateway_action( - str(uuid.uuid4()), hardware_id, ActionMethod.GRIPPER_GET, None + result = self._submit(lambda: self._gripper(False, hardware_id, None)) + if isinstance(result, CommandResult): + return cast("CommandResult[float]", result) + return cast("CommandResult[float]", self._wait_aux(result)) + + def _gripper( + self, setter: bool, hardware: str, position: float | None + ) -> AuxiliaryTicket | CommandResult[Any]: + if self._gate != "open": + return CommandResult(False, diagnostic="gateway is closing", snapshot=self._snapshot()) + if setter and position is None: + return CommandResult( + False, diagnostic="position is required", snapshot=self._snapshot() ) - if not admitted: - return CommandResult( - False, diagnostic="gateway is closing", snapshot=self.snapshot() - ) - if not isinstance(position, (int, float)): - return CommandResult( - False, - diagnostic="gateway returned no gripper position", - snapshot=self.snapshot(), - ) - return CommandResult(True, float(position), snapshot=self.snapshot()) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + action_id = str(uuid.uuid4()) + ticket = self._auxiliary.register(action_id, now + self._action_timeout, setter) + if setter: + assert position is not None + fn: Callable[[], Any] = lambda: self._gateway.set_gripper_position(hardware, position) # noqa: E731 + else: + fn = lambda: self._gateway.get_gripper_position(hardware) # noqa: E731 + self._effects.submit_auxiliary(action_id, fn, self._events.put) + return ticket - return self._submit(fn) + def _wait_aux(self, ticket: AuxiliaryTicket) -> CommandResult[Any]: + start = self._clock_now() + if start is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + end = start + self._action_timeout + with self._condition: + while True: + result = self._auxiliary.take_result(ticket.action_id) + if result is not None: + accepted, value, diagnostic = result + if ticket.setter: + return CommandResult( + accepted and isinstance(value, Outcome), + value if isinstance(value, Outcome) else None, + diagnostic=diagnostic, + snapshot=self._snapshot(), + ) + valid = isinstance(value, (int, float)) and not isinstance(value, bool) + return CommandResult( + accepted and valid, + float(value) if valid else None, + diagnostic=diagnostic, + snapshot=self._snapshot(), + ) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + if now >= end: + return CommandResult( + False, + diagnostic="coordinator auxiliary action timeout", + snapshot=self._snapshot(), + ) + self._condition.wait(max(0.0, end - now)) + + def _wait(self, handle: Any, kind: str, timeout: float) -> CommandResult[Any]: + end = time.monotonic() + max(0.0, timeout) + with self._condition: + while time.monotonic() < end: + values = ( + self._context.dispatch_results + if kind == "dispatch" + else self._context.terminal_results + if kind == "terminal" + else self._context.reset_results + ) + found = dict(values).get(handle) + if kind == "terminal": + fault_diagnostic = self._fault_terminal_results.get(handle) + if fault_diagnostic is not None: + return CommandResult( + False, + diagnostic=fault_diagnostic, + snapshot=self._snapshot(), + ) + if found is not None: + return CommandResult(True, found, snapshot=self._snapshot()) + if kind == "terminal": + active = self._context.active + if ( + active is not None + and active.handle == handle + and self._context.state == LifecycleState.FAULT + ): + return CommandResult( + False, + diagnostic=self._context.diagnostic or "runtime fault", + snapshot=self._snapshot(), + ) + remaining = end - time.monotonic() + self._condition.wait(min(0.1, max(0.0, remaining))) + return CommandResult(False, diagnostic="timeout", snapshot=self._snapshot()) def wait_for_dispatch( self, handle: OperationHandle, timeout: float = 1.0 @@ -1855,230 +1338,137 @@ def wait_for_terminal( ) -> CommandResult[ExecutionResult]: return self._wait(handle, "terminal", timeout) - def _wait( - self, handle: OperationHandle, result_kind: str, timeout: float - ) -> CommandResult[ExecutionResult]: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - with self._condition: - results = ( - self._context.dispatch_results - if result_kind == "dispatch" - else self._context.terminal_results - ) - found = dict(results).get(handle) - if found is None: - self._condition.wait(min(0.01, deadline - time.monotonic())) - if found is not None: - return CommandResult(True, found, snapshot=self.snapshot()) - return CommandResult(False, diagnostic="timeout", snapshot=self.snapshot()) - - def _poll_loop(self) -> None: - while not self._stop.wait(self._poll_interval): - self._events.put(_Call(lambda: self._reduce(_Event("poll_tick")))) - def reset(self) -> CommandResult[ResetHandle]: def fn() -> CommandResult[ResetHandle]: + if self._gate != "open": + return CommandResult( + False, + diagnostic="reset unavailable while runtime is closing", + snapshot=self._snapshot(), + ) + existing = self._context.reset_handle + if existing is not None: + prior = dict(self._context.reset_results).get(existing) + if prior is not None: + return CommandResult( + False, + existing, + diagnostic=prior.diagnostic, + snapshot=self._snapshot(), + ) + return CommandResult(True, existing, snapshot=self._snapshot()) if self._context.state != LifecycleState.FAULT: return CommandResult( - False, diagnostic="reset requires FAULT", snapshot=self.snapshot() + False, diagnostic="reset requires FAULT", snapshot=self._snapshot() ) - handle = self._context.reset_handle or ResetHandle(str(uuid.uuid4())) - if self._context.reset_handle is not None: - return CommandResult(True, self._context.reset_handle, snapshot=self.snapshot()) - deadline = time.monotonic() + max(self._timeout, 0.05) - self._commit( - replace( - self._context, - reset_handle=handle, - reset_deadline=deadline, - reset_proven_inactive=frozenset(), - reset_completed_tasks=frozenset(), + if self._pending: + return CommandResult( + False, + diagnostic="reset blocked by unresolved coordinator RPC", + snapshot=self._snapshot(), ) - ) - self._reconcile_reset(self._context.active) - self._start_auxiliary( - self._reset_deadline, handle, deadline, name=f"reset-deadline-{handle.reset_id}" - ) - return CommandResult(True, handle, snapshot=self.snapshot()) - - return self._submit(fn) - - def _reset_deadline( - self, cancel: threading.Event, handle: ResetHandle, deadline: float - ) -> None: - if cancel.wait(max(0.0, deadline - time.monotonic())): - return - self._events.put( - _Call( - lambda: self._reduce( - _Event( - "reset_deadline", - payload=handle.reset_id, - reset_id=handle.reset_id, - ) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() ) + h = ResetHandle(str(uuid.uuid4())) + self._commit( + reset_handle=h, + reset_deadline=now + self._action_timeout, + reset_proven_inactive=frozenset(), + reset_completed_tasks=frozenset(), ) - ) + op = self._context.active + if op: + self._reset_next(op) + else: + self._finish_reset(True) + return CommandResult(True, h, snapshot=self._snapshot()) + + return self._submit(fn, closing=True) def wait_for_reset( self, handle: ResetHandle, timeout: float = 1.0 ) -> CommandResult[ResetResult]: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - with self._condition: - result = dict(self._context.reset_results).get(handle) - if result is None: - self._condition.wait(min(0.01, deadline - time.monotonic())) - if result is not None: - return CommandResult(True, result, snapshot=self.snapshot()) - return CommandResult(False, diagnostic="timeout", snapshot=self.snapshot()) - - def _seal_gateway(self) -> None: - with self._gateway_condition: - if self._gateway_gate in (_GatewayGateState.SEALED, _GatewayGateState.STOPPED): - return - self._gateway_gate = _GatewayGateState.SEALED - self._gateway_condition.notify_all() - self._start_auxiliary(self._close_gateway, name="gateway-close") - - def _close_gateway(self, _cancel: threading.Event) -> None: - with self._gateway_condition: - while self._rpc_inflight: - self._gateway_condition.wait() - if self._gateway_close_started: - return - self._gateway_close_started = True - try: - self._gateway.stop() - except BaseException as exc: - with self._gateway_condition: - self._gateway_close_error = f"gateway close failed: {exc}" - self._gateway_condition.notify_all() - return - else: - with self._gateway_condition: - self._gateway_gate = _GatewayGateState.STOPPED - self._gateway_close_done = True - self._gateway_condition.notify_all() - - def _wait_gateway_closed(self, deadline: float) -> bool: - with self._gateway_condition: - while ( - not self._gateway_close_done - and self._gateway_close_error is None - and time.monotonic() < deadline - ): - self._gateway_condition.wait(min(0.01, deadline - time.monotonic())) - return self._gateway_close_done and self._gateway_close_error is None + return self._wait(handle, "reset", timeout) def shutdown(self, timeout: float = 1.0) -> CommandResult[ShutdownResult]: with self._shutdown_lock: - if self._closed: - result = self._context.shutdown_result or ShutdownResult( - False, "runtime is closed without a shutdown result" - ) - return CommandResult(result.success, result, snapshot=self.snapshot()) - if not self._owner.is_alive() and self._context.shutdown == ShutdownState.CLOSING: - result = self._context.shutdown_result or ShutdownResult( - False, "shutdown result unavailable" - ) - if not result.success: - return CommandResult(False, result, snapshot=self.snapshot()) - self._commit( - replace(self._context, shutdown=ShutdownState.CLOSED, shutdown_result=result) - ) - self._closed = True - return CommandResult(result.success, result, snapshot=self.snapshot()) - deadline = time.monotonic() + max(0.0, timeout) - with self._condition: - with self._gateway_condition: - self._closing_requested = True - self._context = replace( - self._context, - shutdown=ShutdownState.CLOSING, - shutdown_deadline=deadline, - revision=self._context.revision + 1, - ) - if self._gateway_gate == _GatewayGateState.OPEN: - self._gateway_gate = _GatewayGateState.CANCEL_ONLY - self._gateway_condition.notify_all() - self._condition.notify_all() - - self._stop.set() - call = _Call(lambda: self._reduce(_Event("shutdown", payload=deadline))) - self._events.put(call) - remaining = max(0.0, deadline - time.monotonic()) - try: - call.answer.get(timeout=remaining) - except queue.Empty: - pass - - with self._condition: - while self._context.shutdown_result is None and time.monotonic() < deadline: - self._condition.wait(min(0.01, deadline - time.monotonic())) - - if self._context.shutdown_result is None: - deadline_call = _Call( - lambda: self._reduce(_Event("shutdown_deadline", payload=deadline)) - ) - self._events.put(deadline_call) - remaining = max(0.0, deadline - time.monotonic()) - if remaining: - try: - deadline_call.answer.get(timeout=remaining) - except queue.Empty: - pass - if self._context.shutdown_result is None: - self._commit( - replace( - self._context, - shutdown_result=ShutdownResult(False, "shutdown safety deadline exceeded"), - ) - ) - - self._seal_gateway() - close_done = self._wait_gateway_closed(deadline) - if not close_done and ( - self._context.shutdown_result is None or self._context.shutdown_result.success - ): - with self._gateway_condition: - close_diagnostic = ( - self._gateway_close_error - or "gateway close unresolved before shutdown deadline" + with self._lock: + if self._context.shutdown_result is not None: + return CommandResult( + self._context.shutdown_result.success, + self._context.shutdown_result, + snapshot=self._snapshot(), ) - self._commit( - replace(self._context, shutdown_result=ShutdownResult(False, close_diagnostic)) - ) - auxiliary_drained = self._stop_auxiliary(deadline) - if not auxiliary_drained and ( - self._context.shutdown_result is None or self._context.shutdown_result.success - ): - self._commit( - replace( - self._context, - shutdown_result=ShutdownResult( - False, "auxiliary runtime thread unresolved before shutdown deadline" - ), + if self._shutdown_initiated: + return CommandResult( + False, + diagnostic="shutdown already initiated", + snapshot=self._snapshot(), ) + self._shutdown_initiated = True + self._gate = "cancel_only" + return self._shutdown_wait(timeout) + + def _shutdown_wait(self, timeout: float) -> CommandResult[ShutdownResult]: + start = self._clock_now() + if start is None: + end = time.monotonic() + 1.0 + with self._condition: + while self._context.shutdown_result is None and time.monotonic() < end: + self._condition.wait(0.01) + result = self._context.shutdown_result + if result is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self.snapshot() ) - self._events.put(None) - remaining = max(0.0, deadline - time.monotonic()) - self._owner.join(remaining) - owner_stopped = not self._owner.is_alive() - result = self._context.shutdown_result or ShutdownResult( - False, "shutdown result unavailable" + return CommandResult(False, result, snapshot=self.snapshot()) + deadline = start + max(0.0, timeout) + self._submit(lambda: self._shutdown_owner(deadline), closing=True) + while True: + with self._condition: + if self._context.shutdown_result is not None: + break + now = self._clock_now() + if now is None or now >= deadline: + break + with self._condition: + self._condition.wait(min(0.1, max(0.0, deadline - now))) + if self._context.shutdown_result is None: + self._submit( + lambda: self._finish_shutdown(False, "shutdown safety deadline exceeded"), + closing=True, ) - if not owner_stopped: - result = ShutdownResult(False, "owner thread did not stop before shutdown deadline") - self._commit(replace(self._context, shutdown_result=result)) - return CommandResult(False, result, snapshot=self.snapshot()) - if not result.success: - return CommandResult(False, result, snapshot=self.snapshot()) - self._commit( - replace(self._context, shutdown=ShutdownState.CLOSED, shutdown_result=result) + result = self._context.shutdown_result + if result is None: + return CommandResult( + False, + diagnostic="shutdown safety deadline exceeded", + snapshot=self.snapshot(), + ) + if not result.success: + return CommandResult(False, result, snapshot=self.snapshot()) + self._owner.join() + if self._owner.is_alive(): + return CommandResult( + False, + diagnostic="owner thread did not stop", + snapshot=self.snapshot(), ) - self._closed = True - return CommandResult(result.success and owner_stopped, result, snapshot=self.snapshot()) + self._closed = True + return CommandResult(True, result, snapshot=self.snapshot()) close = shutdown + + def _shutdown_owner(self, deadline: float) -> CommandResult[None]: + self._commit( + shutdown=ShutdownState.CLOSING, + shutdown_deadline=deadline, + state=LifecycleState.CANCELLING if self._context.active else LifecycleState.IDLE, + ) + if self._context.active: + self._begin_cancellation("shutdown cancellation requested") + self._shutdown_progress() + return CommandResult(True, snapshot=self._snapshot()) diff --git a/dimos/manipulation/execution_topology.py b/dimos/manipulation/execution_topology.py new file mode 100644 index 0000000000..d3a528c30a --- /dev/null +++ b/dimos/manipulation/execution_topology.py @@ -0,0 +1,273 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Pure topology binding and generated-plan materialization helpers.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +@dataclass(frozen=True) +class ExecutionTopology: + """Canonical group -> robot -> coordinator-task routing.""" + + routes: tuple[tuple[str, str, str], ...] + robot_configs: tuple[RobotModelConfig, ...] = () + inverse_joint_mappings: tuple[tuple[str, tuple[tuple[str, str], ...]], ...] = () + + @classmethod + def from_robot_configs( + cls, + robot_configs: Iterable[RobotModelConfig], + group_routes: Mapping[str, tuple[str, str]] | Iterable[tuple[str, str, str]] | None = None, + group_resolver: Callable[[str], tuple[str, str]] | None = None, + ) -> ExecutionTopology: + """Build canonical routes while retaining the supplied config references.""" + configs = tuple(robot_configs) + by_robot: dict[str, RobotModelConfig] = {} + for config in configs: + if config.name in by_robot: + raise ValueError(f"duplicate robot config: {config.name}") + by_robot[config.name] = config + explicit: dict[str, tuple[str, str]] = {} + if isinstance(group_routes, Mapping): + for group, route in group_routes.items(): + if not isinstance(group, str) or not isinstance(route, tuple) or len(route) != 2: + raise ValueError("invalid group route mapping") + explicit[group] = (str(route[0]), str(route[1])) + elif group_routes is not None: + for group, robot, task in group_routes: + if group in explicit: + raise ValueError(f"duplicate group route: {group}") + explicit[group] = (robot, task) + routes: list[tuple[str, str, str]] = [] + inverse: list[tuple[str, tuple[tuple[str, str], ...]]] = [] + for config in configs: + local_names = set(config.joint_names) + local_to_coordinator: dict[str, str] = {} + for coordinator_name, local_name in config.joint_name_mapping.items(): + if local_name not in local_names: + raise ValueError( + f"mapped local joint is not configured: {config.name}/{local_name}" + ) + previous = local_to_coordinator.get(local_name) + if previous is not None and previous != coordinator_name: + raise ValueError( + f"duplicate local joint mapping for {config.name}/{local_name}" + ) + local_to_coordinator[local_name] = coordinator_name + coordinator_names: set[str] = set() + for local_name in config.joint_names: + coordinator_name = local_to_coordinator.get(local_name, local_name) + if coordinator_name in coordinator_names: + raise ValueError( + f"duplicate coordinator joint mapping for {config.name}/{coordinator_name}" + ) + coordinator_names.add(coordinator_name) + inverse.append( + ( + config.name, + tuple( + (local, local_to_coordinator.get(local, local)) + for local in config.joint_names + ), + ) + ) + for definition in config.planning_groups: + if not set(definition.joint_names).issubset(local_names): + raise ValueError( + f"planning group joints are not configured: {config.name}/{definition.name}" + ) + group_id = f"{config.name}/{definition.name}" + if group_resolver is not None: + route_robot, route_task = group_resolver(group_id) + else: + if not config.coordinator_task_name and group_id not in explicit: + continue + route_robot, route_task = explicit.get( + group_id, (config.name, config.coordinator_task_name or "") + ) + if route_robot != config.name or route_task != config.coordinator_task_name: + raise ValueError(f"route/config mismatch for {group_id}") + routes.append((group_id, route_robot, route_task)) + if explicit and set(explicit) != {group for group, _, _ in routes}: + raise ValueError("route set does not match configured planning groups") + return cls(tuple(routes), configs, tuple(inverse)) + + from_configs = from_robot_configs + from_robot_model_configs = from_robot_configs + + def config_for_robot(self, robot_name: str) -> RobotModelConfig: + for config in self.robot_configs: + if config.name == robot_name: + return config + raise KeyError(f"unknown robot: {robot_name}") + + def coordinator_joint_name(self, robot_name: str, local_name: str) -> str: + for name, mappings in self.inverse_joint_mappings: + if name == robot_name: + for local, coordinator in mappings: + if local == local_name: + return coordinator + break + return local_name + + def route_for_group(self, group_id: str) -> tuple[str, str]: + routes = {group: (robot, task) for group, robot, task in self.routes} + try: + return routes[group_id] + except KeyError as exc: + raise ValueError(f"missing route for planning group {group_id}") from exc + + def affected_robots(self, group_ids: Iterable[str]) -> tuple[str, ...]: + robots: list[str] = [] + for group_id in group_ids: + robot_name, _task_name = self.route_for_group(group_id) + if robot_name in robots: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + robots.append(robot_name) + return tuple(robots) + + @property + def groups(self) -> tuple[str, ...]: + return tuple(group for group, _, _ in self.routes) + + +@dataclass(frozen=True) +class TaskEntry: + planning_group: str + task_name: str + request: Any + robot_name: str | None = None + + +@dataclass(frozen=True) +class ExecutionPlan: + generated_plan: Any + entries: tuple[TaskEntry, ...] + + +@dataclass(frozen=True) +class PreparedPlan: + generated_plan: GeneratedPlan + entries: tuple[TaskEntry, ...] + topology: ExecutionTopology | None = None + _canonical: bool = field(default=False, repr=False, compare=False) + + +def prepare_generated_plan( + generated_plan: GeneratedPlan, topology: ExecutionTopology +) -> PreparedPlan: + """Purely bind a generated multi-robot plan to canonical coordinator tasks.""" + if not isinstance(generated_plan, GeneratedPlan): + raise TypeError("expected GeneratedPlan") + group_ids = tuple(generated_plan.group_ids) + if not group_ids or len(set(group_ids)) != len(group_ids): + raise ValueError("generated plan must contain unique planning groups") + route_by_group = {group: (robot, task) for group, robot, task in topology.routes} + if len(route_by_group) != len(topology.routes): + raise ValueError("topology contains duplicate group routes") + selected_routes: list[tuple[str, str, str, RobotModelConfig, tuple[str, ...]]] = [] + affected: set[str] = set() + for group_id in group_ids: + route = route_by_group.get(group_id) + if route is None: + robot_name = group_id.split("/", 1)[0] + try: + if not topology.config_for_robot(robot_name).coordinator_task_name: + raise ValueError(f"no coordinator task for selected robot {robot_name}") + except KeyError: + pass + raise ValueError(f"missing route for planning group {group_id}") + robot_name, task_name = route + if not task_name: + raise ValueError(f"missing coordinator task for {group_id}") + if robot_name in affected: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + config = topology.config_for_robot(robot_name) + prefix, _, definition_name = group_id.partition("/") + if prefix != robot_name or not definition_name: + raise ValueError(f"invalid canonical planning group {group_id}") + definitions = [ + definition + for definition in config.planning_groups + if definition.name == definition_name + ] + if len(definitions) != 1: + raise ValueError(f"missing config group for {group_id}") + expected_local = tuple(definitions[0].joint_names) + if len(set(expected_local)) != len(expected_local): + raise ValueError(f"duplicate intended joints for {group_id}") + affected.add(robot_name) + selected_routes.append((group_id, robot_name, task_name, config, expected_local)) + expected_by_robot = {robot: set(expected) for _, robot, _, _, expected in selected_routes} + + trajectory = generated_plan.trajectory + names = tuple(trajectory.joint_names) + if not names: + raise ValueError("generated trajectory has no joint names") + if len(set(names)) != len(names): + raise ValueError("duplicate global trajectory joint assignment") + columns: dict[str, dict[str, int]] = {robot: {} for robot in affected} + known_robots = {config.name for config in topology.robot_configs} + for index, global_name in enumerate(names): + if global_name.count("/") != 1: + raise ValueError(f"malformed global joint name: {global_name}") + robot_name, local_name = global_name.split("/", 1) + if not robot_name or not local_name or robot_name not in known_robots: + raise ValueError(f"unknown global joint assignment: {global_name}") + if robot_name in columns: + if local_name not in expected_by_robot[robot_name]: + raise ValueError(f"additional joint column for selected robot: {global_name}") + if local_name in columns[robot_name]: + raise ValueError(f"duplicate robot joint assignment: {global_name}") + columns[robot_name][local_name] = index + points = tuple(trajectory.points) + for point in points: + if len(point.positions) != len(names): + raise ValueError("trajectory point dimension does not match joint names") + if point.velocities and len(point.velocities) != len(names): + raise ValueError("trajectory velocity dimension does not match joint names") + + entries: list[TaskEntry] = [] + for group_id, robot_name, task_name, _config, expected_local in selected_routes: + robot_columns = columns[robot_name] + missing = [local for local in expected_local if local not in robot_columns] + if missing: + raise ValueError(f"missing robot joints for {robot_name}: {missing}") + intended = set(expected_local) + local_order = tuple(local for local in robot_columns if local in intended) + coordinator_names = tuple( + topology.coordinator_joint_name(robot_name, local) for local in local_order + ) + if len(set(coordinator_names)) != len(coordinator_names): + raise ValueError(f"duplicate coordinator joint names for {robot_name}") + local_trajectory = JointTrajectory( + joint_names=list(coordinator_names), + points=[ + TrajectoryPoint( + time_from_start=point.time_from_start, + positions=[point.positions[robot_columns[local]] for local in local_order], + velocities=( + [point.velocities[robot_columns[local]] for local in local_order] + if point.velocities + else None + ), + ) + for point in points + ], + ) + entries.append(TaskEntry(group_id, task_name, {"trajectory": local_trajectory}, robot_name)) + return PreparedPlan(generated_plan, tuple(entries), topology, True) + + +prepare_plan = prepare_generated_plan +prepare_execution_plan = prepare_generated_plan +materialize_prepared_plan = prepare_generated_plan diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index b4a157765e..c80b31e072 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -47,6 +47,7 @@ OperationHandle, Outcome, PreparedPlan, + RuntimeSnapshot, ShutdownState, prepare_generated_plan, ) @@ -128,6 +129,7 @@ class ManipulationModuleConfig(ModuleConfig): robots: list[RobotModelConfig] = Field(default_factory=list) planning_timeout: float = 10.0 + physical_operation_timeout: float = Field(default=60.0, gt=0, allow_inf_nan=False) world_backend: WorldBackend = "drake" visualization: ManipulationVisualizationConfig = Field( default_factory=NoManipulationVisualizationConfig @@ -285,6 +287,7 @@ def gateway_factory() -> ControlCoordinatorGateway: gateway_factory, topology=self._execution_topology, action_timeout=self.config.planning_timeout, + physical_operation_timeout=self.config.physical_operation_timeout, ) # Add floor obstacle to prevent trajectories below the table surface @@ -445,21 +448,27 @@ def _tf_publish_loop(self) -> None: @rpc def get_state(self) -> str: """Get current manipulation state name.""" - if self._execution_runtime is None: + snapshot = self._execution_snapshot() + if snapshot is None: return LifecycleState.IDLE.name - return self._execution_runtime.snapshot().state.name + return snapshot.state.name + + def _execution_snapshot(self) -> RuntimeSnapshot | None: + """Read the runtime-owned execution projection exactly once.""" + runtime = self._execution_runtime + return runtime.snapshot() if runtime is not None else None @rpc def get_execution_snapshot(self) -> ManipulationExecutionSnapshot: """Return one atomic runtime execution snapshot for observers.""" - if self._execution_runtime is None: + snapshot = self._execution_snapshot() + if snapshot is None: return ManipulationExecutionSnapshot( state=LifecycleState.IDLE.name, diagnostic="", ready_plan_id=None, has_ready_plan=False, ) - snapshot = self._execution_runtime.snapshot() return ManipulationExecutionSnapshot( state=snapshot.state.name, diagnostic=snapshot.diagnostic or "", @@ -474,9 +483,9 @@ def get_error(self) -> str: Returns: Error message or empty string """ - if self._execution_runtime is None: + snapshot = self._execution_snapshot() + if snapshot is None: return "" - snapshot = self._execution_runtime.snapshot() return snapshot.diagnostic or "" @rpc @@ -530,11 +539,8 @@ def reset(self) -> SkillResult[ManipulationSkillError]: reset_result = self._execution_runtime.reset() if not reset_result.accepted or reset_result.value is None: return SkillResult.fail("INVALID_STATE", reset_result.diagnostic) - # Runtime owns the reset deadline (action timeout); wait slightly past - # it so the deadline event has a chance to publish its result. - runtime_timeout = float(getattr(self._execution_runtime, "_timeout", 0.05)) reconciled = self._execution_runtime.wait_for_reset( - reset_result.value, timeout=max(runtime_timeout, 0.05) + 0.1 + reset_result.value, timeout=float("inf") ) if not reconciled.accepted or reconciled.value is None or not reconciled.value.success: return SkillResult.fail("RESET_FAILED", reconciled.diagnostic or "Reset failed") @@ -1111,9 +1117,10 @@ def has_planned_path(self) -> bool: Returns: True if a path is planned and ready """ - if self._execution_runtime is None: + snapshot = self._execution_snapshot() + if snapshot is None: return False - ready = self._execution_runtime.snapshot().ready_plan + ready = snapshot.ready_plan return isinstance(ready, PreparedPlan) and bool(ready.generated_plan.path) @rpc @@ -1371,21 +1378,15 @@ def set_init_joints_to_current(self, robot_name: RobotName | None = None) -> boo return True @rpc - def execute(self) -> bool: - """Execute the current runtime-owned prepared plan.""" - handle = self._submit_execution(None) + def execute(self, plan: GeneratedPlan | None = None) -> bool: + """Execute a complete plan, or the runtime-owned READY plan when omitted.""" + handle = self._submit_execution(plan) if handle is None or self._execution_runtime is None: return False - dispatch = self._execution_runtime.wait_for_dispatch( - handle, timeout=self._dispatch_wait_timeout() - ) - if not dispatch.accepted: - return self._handle_dispatch_timeout(handle, dispatch.diagnostic) == "completed" - return bool( - dispatch.accepted - and dispatch.value is not None - and dispatch.value.outcome == Outcome.ACCEPTED - ) + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: + return False + return bool(dispatch.value.outcome == Outcome.ACCEPTED) @rpc def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: @@ -1393,16 +1394,10 @@ def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: handle = self._submit_execution(plan) if handle is None or self._execution_runtime is None: return False - dispatch = self._execution_runtime.wait_for_dispatch( - handle, timeout=self._dispatch_wait_timeout() - ) - if not dispatch.accepted: - return self._handle_dispatch_timeout(handle, dispatch.diagnostic) == "completed" - return bool( - dispatch.accepted - and dispatch.value is not None - and dispatch.value.outcome == Outcome.ACCEPTED - ) + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: + return False + return bool(dispatch.value.outcome == Outcome.ACCEPTED) def _submit_execution(self, plan: GeneratedPlan | None) -> OperationHandle | None: """Submit one explicit or exactly-ready generated plan to the runtime.""" @@ -1410,189 +1405,23 @@ def _submit_execution(self, plan: GeneratedPlan | None) -> OperationHandle | Non topology = self._execution_topology if runtime is None or topology is None: return None - ready = runtime.snapshot().ready_plan - if plan is None or (isinstance(ready, PreparedPlan) and ready.generated_plan is plan): + if plan is None: result = runtime.execute_ready() else: - previous = ready.generated_plan if isinstance(ready, PreparedPlan) else None try: prepared = prepare_generated_plan(plan, topology) except Exception as exc: logger.warning("Failed to prepare execution: %s", exc) return None result = runtime.execute_explicit(prepared) - if result.accepted and previous is not None: - self._dismiss_preview(previous.group_ids) return result.value if result.accepted else None - def _dispatch_wait_timeout(self) -> float: - """Return the runtime dispatch deadline plus a publication margin.""" - runtime = self._execution_runtime - configured = self.config.planning_timeout - action_timeout = float(getattr(runtime, "_timeout", configured)) if runtime else configured - return max(action_timeout, 0.05) + 0.1 - - def _handle_dispatch_timeout(self, handle: OperationHandle, diagnostic: str) -> str: - """Cancel the timed-out operation and verify its terminal safety state.""" - runtime = self._execution_runtime - if runtime is None: - logger.error("Execution dispatch timed out; runtime is unavailable") - return "unresolved" - - # A timeout can race with result publication. Reconcile the original - # handle before inspecting the active operation, so a replacement can - # never be cancelled on behalf of this request. - retained_dispatch = runtime.wait_for_dispatch(handle, timeout=0.01) - retained_terminal = runtime.wait_for_terminal(handle, timeout=0.01) - terminal_result = ( - retained_terminal.value - if retained_terminal.accepted and retained_terminal.value is not None - else None - ) - if terminal_result is not None: - if terminal_result.outcome == Outcome.CANCELLED: - logger.warning("Execution dispatch timed out; operation was safely cancelled") - return "cancelled" - if terminal_result.outcome == Outcome.COMPLETED: - logger.info("Execution dispatch timed out; operation completed") - return "completed" - if terminal_result.outcome == Outcome.REJECTED: - logger.error("Execution dispatch timed out; execution was rejected") - return "rejected" - if terminal_result.outcome == Outcome.FAILED: - logger.error("Execution dispatch timed out; execution failed") - return "failed" - logger.error( - "Execution dispatch timed out; remote safety is unresolved (%s)", - terminal_result.outcome.value, - ) - return "unresolved" - if ( - retained_dispatch.accepted - and retained_dispatch.value is not None - and retained_dispatch.value.outcome != Outcome.ACCEPTED - ): - if retained_dispatch.value.outcome == Outcome.CANCELLED: - logger.warning("Execution dispatch timed out; operation was safely cancelled") - return "cancelled" - if retained_dispatch.value.outcome == Outcome.COMPLETED: - logger.info("Execution dispatch timed out; operation completed") - return "completed" - if retained_dispatch.value.outcome == Outcome.REJECTED: - logger.error("Execution dispatch timed out; execution was rejected") - return "rejected" - if retained_dispatch.value.outcome == Outcome.FAILED: - logger.error("Execution dispatch timed out; execution failed") - return "failed" - logger.error( - "Execution dispatch timed out; remote safety is unresolved (%s)", - retained_dispatch.value.outcome.value, - ) - return "unresolved" - - cancellation = runtime.cancel_if_current(handle) - if not cancellation.accepted: - concurrent_terminal = runtime.wait_for_terminal(handle, timeout=0.01) - if concurrent_terminal.accepted and concurrent_terminal.value is not None: - outcome = concurrent_terminal.value.outcome - if outcome == Outcome.CANCELLED: - logger.warning("Execution dispatch timed out; operation was safely cancelled") - return "cancelled" - if outcome == Outcome.COMPLETED: - logger.info("Execution dispatch timed out; operation completed") - return "completed" - if outcome == Outcome.REJECTED: - logger.error("Execution dispatch timed out; execution was rejected") - return "rejected" - if outcome == Outcome.FAILED: - logger.error("Execution dispatch timed out; execution failed") - return "failed" - logger.error( - "Execution dispatch timed out; remote safety is unresolved (%s)", - outcome.value, - ) - return "unresolved" - active = runtime.snapshot().operation - if active is None or active.handle != handle: - message = ( - "Execution dispatch timed out; original operation is unresolved and " - "active operation handle changed" - if active is not None - else "Execution dispatch timed out; original operation is unresolved and no active operation remains" - ) - logger.error("%s (%s)", message, diagnostic or "timeout") - return "unresolved" - message = ( - cancellation.diagnostic or "Execution dispatch timed out; cancellation was rejected" - ) - logger.error("%s", message) - return "cancel_rejected" - - terminal = runtime.wait_for_terminal(handle, timeout=self._dispatch_wait_timeout()) - if terminal.accepted and terminal.value is not None: - if terminal.value.outcome == Outcome.CANCELLED: - logger.warning("Execution dispatch timed out; operation was safely cancelled") - return "cancelled" - if terminal.value.outcome == Outcome.COMPLETED: - logger.info("Execution dispatch timed out; operation completed") - return "completed" - if terminal.value.outcome == Outcome.REJECTED: - logger.error("Execution dispatch timed out; execution was rejected") - return "rejected" - if terminal.value.outcome == Outcome.FAILED: - logger.error("Execution dispatch timed out; execution failed") - return "failed" - logger.error( - "Execution dispatch timed out; remote safety is unresolved (%s)", - terminal.value.outcome.value, - ) - return "unresolved" - else: - message = ( - "Execution dispatch timed out; cancellation was accepted but operation " - "remains unresolved" - ) - logger.error( - "%s%s", - message, - f" ({terminal.diagnostic})" if terminal.diagnostic else "", - ) - return "unresolved" - - def _terminal_after_cancel_rejection( - self, handle: OperationHandle, diagnostic: str - ) -> SkillResult[ManipulationSkillError] | None: - """Reconcile a cancellation race using only the original handle.""" - runtime = self._execution_runtime - if runtime is None: - return None - terminal = runtime.wait_for_terminal(handle, timeout=0.01) - if not terminal.accepted or terminal.value is None: - return None - result = terminal.value - if result.outcome == Outcome.COMPLETED: - return SkillResult.ok("Execution completed while cancellation was requested") - if result.outcome == Outcome.CANCELLED: - return SkillResult.fail( - "EXECUTION_CANCELLED", result.diagnostic or "Trajectory cancelled" - ) - if result.outcome == Outcome.REJECTED: - return SkillResult.fail( - "EXECUTION_REJECTED", result.diagnostic or "Trajectory rejected" - ) - if result.outcome == Outcome.FAILED: - return SkillResult.fail("EXECUTION_FAILED", result.diagnostic or "Trajectory failed") - return SkillResult.fail( - "EXECUTION_CANCEL_UNRESOLVED", - result.diagnostic or diagnostic or "Remote execution safety is unresolved", - ) - @rpc def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: """Return a snapshot-only runtime status, optionally scoped to a robot.""" - if self._execution_runtime is None: + snapshot = self._execution_snapshot() + if snapshot is None: return None - snapshot = self._execution_runtime.snapshot() operation = snapshot.operation tasks: list[dict[str, Any]] = [] if operation is not None: @@ -1789,40 +1618,11 @@ def _preview_execute_wait( if handle is None: return SkillResult.fail("EXECUTION_FAILED", "Trajectory execution failed") - dispatch = self._execution_runtime.wait_for_dispatch( - handle, timeout=self._dispatch_wait_timeout() - ) - if not dispatch.accepted: - safety = self._handle_dispatch_timeout(handle, dispatch.diagnostic) - if safety == "cancel_rejected": - return SkillResult.fail( - "EXECUTION_CANCEL_REJECTED", - dispatch.diagnostic - or "Execution dispatch timed out; cancellation was rejected", - ) - if safety == "cancelled": - return SkillResult.fail( - "EXECUTION_CANCELLED", - "Execution dispatch timed out; operation was safely cancelled", - ) - if safety == "completed": - return SkillResult.ok("Execution dispatch timed out; operation completed") - if safety == "failed": - return SkillResult.fail( - "EXECUTION_FAILED", - "Execution dispatch timed out; operation resolved with failure", - ) - if safety == "rejected": - return SkillResult.fail( - "EXECUTION_REJECTED", "Execution dispatch timed out; execution was rejected" - ) - return SkillResult.fail( - "EXECUTION_CANCEL_UNRESOLVED", - "Execution dispatch timed out; cancellation was accepted but operation is unresolved", - ) - if dispatch.value is None: + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: return SkillResult.fail( - "EXECUTION_DISPATCH_FAILED", "Execution dispatch returned no result" + "EXECUTION_DISPATCH_FAILED", + dispatch.diagnostic or "Execution dispatch failed", ) if dispatch.value.outcome == Outcome.REJECTED: return SkillResult.fail( @@ -1831,59 +1631,29 @@ def _preview_execute_wait( if dispatch.value.outcome != Outcome.ACCEPTED: return SkillResult.fail( "EXECUTION_DISPATCH_FAILED", - dispatch.value.diagnostic - or f"Unexpected dispatch outcome: {dispatch.value.outcome.value}", - ) - - duration = plan.trajectory.duration if plan is not None else 0.0 - execution_timeout = max(float(duration) + 1.0, 0.5) - terminal = self._execution_runtime.wait_for_terminal(handle, timeout=execution_timeout) - if not terminal.accepted: - cancellation = self._execution_runtime.cancel_if_current(handle) - if not cancellation.accepted: - reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) - if reconciled is not None: - return reconciled - return SkillResult.fail( - "EXECUTION_CANCEL_REJECTED", - cancellation.diagnostic or "Execution timed out and cancellation was rejected", - ) - return SkillResult.fail( - "EXECUTION_CANCEL_UNRESOLVED", - "Execution timed out; cancellation was accepted but motion is unresolved", + dispatch.value.diagnostic or "Unexpected dispatch outcome", ) - if terminal.value is None: - cancellation = self._execution_runtime.cancel_if_current(handle) - if not cancellation.accepted: - reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) - if reconciled is not None: - return reconciled - return SkillResult.fail("EXECUTION_FAILED", "Execution returned no result") - outcome = terminal.value.outcome - if outcome == Outcome.REJECTED: - return SkillResult.fail( - "EXECUTION_REJECTED", terminal.value.diagnostic or "Trajectory rejected" - ) - if outcome == Outcome.FAILED: + terminal = self._execution_runtime.wait_for_terminal(handle, timeout=float("inf")) + if not terminal.accepted or terminal.value is None: return SkillResult.fail( - "EXECUTION_FAILED", terminal.value.diagnostic or "Trajectory failed" + "EXECUTION_FAILED", terminal.diagnostic or "Execution terminal result unavailable" ) - if outcome == Outcome.CANCELLED: + result = terminal.value + if result.outcome == Outcome.COMPLETED: + return SkillResult.ok() + if result.outcome == Outcome.CANCELLED: return SkillResult.fail( - "EXECUTION_CANCELLED", terminal.value.diagnostic or "Trajectory cancelled" + "EXECUTION_CANCELLED", result.diagnostic or "Trajectory cancelled" ) - if outcome != Outcome.COMPLETED: - cancellation = self._execution_runtime.cancel_if_current(handle) - if not cancellation.accepted: - reconciled = self._terminal_after_cancel_rejection(handle, cancellation.diagnostic) - if reconciled is not None: - return reconciled + if result.outcome == Outcome.FAILED: + return SkillResult.fail("EXECUTION_FAILED", result.diagnostic or "Trajectory failed") + if result.outcome == Outcome.REJECTED: return SkillResult.fail( - "EXECUTION_CANCEL_UNRESOLVED", - f"Remote execution safety is unresolved: {outcome.value}", + "EXECUTION_REJECTED", result.diagnostic or "Trajectory rejected" ) - - return SkillResult.ok() + return SkillResult.fail( + "EXECUTION_FAILED", result.diagnostic or "Execution terminal result was unresolved" + ) @skill def get_robot_state(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: diff --git a/dimos/manipulation/test_execution_auxiliary.py b/dimos/manipulation/test_execution_auxiliary.py new file mode 100644 index 0000000000..bc190af221 --- /dev/null +++ b/dimos/manipulation/test_execution_auxiliary.py @@ -0,0 +1,13 @@ +from dimos.manipulation.execution_auxiliary import AuxiliaryCallBook +from dimos.manipulation.execution_effects import AuxiliaryDone + + +def test_expire_retains_timeout_result_through_late_completion() -> None: + book = AuxiliaryCallBook() + ticket = book.register("aux", 1.0, setter=True) + + assert book.expire(2.0, "timeout") + assert book.has_inflight() + assert book.complete(AuxiliaryDone(ticket.action_id, "late value")) + assert not book.has_unsettled() + assert book.take_result(ticket.action_id) == (False, None, "timeout") diff --git a/dimos/manipulation/test_execution_clock.py b/dimos/manipulation/test_execution_clock.py new file mode 100644 index 0000000000..79a841cb69 --- /dev/null +++ b/dimos/manipulation/test_execution_clock.py @@ -0,0 +1,82 @@ +"""Tests for the serialized, validated execution clock.""" + +import threading + +import pytest + +from dimos.manipulation.execution_clock import InvalidMonotonicClock, ValidatedMonotonicClock + + +class GatedSource: + def __init__(self, initial: float) -> None: + self.first_entry = threading.Event() + self.second_entry = threading.Event() + self.first_release = threading.Event() + self.second_release = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + self.value = initial + + def __call__(self) -> float: + with self._lock: + self.calls += 1 + call = self.calls + if call == 2: + self.first_entry.set() + if not self.first_release.wait(1): + raise AssertionError("first clock sample was not released") + elif call == 3: + self.second_entry.set() + if not self.second_release.wait(1): + raise AssertionError("second clock sample was not released") + return self.value + + +def test_now_serializes_source_sampling_before_validation() -> None: + source = GatedSource(0.0) + clock = ValidatedMonotonicClock(source) + first_done = threading.Event() + second_done = threading.Event() + + def first() -> None: + clock.now() + first_done.set() + + def second() -> None: + clock.now() + second_done.set() + + one = threading.Thread(target=first) + two = threading.Thread(target=second) + one.start() + assert source.first_entry.wait(1) + two.start() + assert not source.second_entry.wait(0.02) + source.first_release.set() + assert first_done.wait(1) + assert source.second_entry.wait(1) + source.second_release.set() + assert second_done.wait(1) + one.join(1) + two.join(1) + + +@pytest.mark.parametrize("sample", [float("nan"), float("inf"), float("-inf")]) +def test_clock_rejects_nonfinite_initial_sample(sample: float) -> None: + with pytest.raises(InvalidMonotonicClock): + ValidatedMonotonicClock(lambda: sample) + + +@pytest.mark.parametrize("sample", [float("nan"), float("inf"), float("-inf")]) +def test_clock_rejects_nonfinite_later_sample(sample: float) -> None: + values = iter((0.0, sample)) + clock = ValidatedMonotonicClock(lambda: next(values)) + with pytest.raises(InvalidMonotonicClock): + clock.now() + + +def test_clock_rejects_decreasing_sample() -> None: + values = iter((1.0, 0.5)) + clock = ValidatedMonotonicClock(lambda: next(values)) + with pytest.raises(InvalidMonotonicClock): + clock.now() diff --git a/dimos/manipulation/test_execution_effects.py b/dimos/manipulation/test_execution_effects.py new file mode 100644 index 0000000000..b6afd2c0b6 --- /dev/null +++ b/dimos/manipulation/test_execution_effects.py @@ -0,0 +1,91 @@ +"""Direct coverage for the daemonized execution-effect runner.""" + +from dataclasses import FrozenInstanceError +import subprocess +import sys +import threading + +import pytest + +from dimos.manipulation.execution_effects import ( + AuxiliaryDone, + EffectDone, + ExecutionEffectRunner, + StopDone, +) + + +def test_four_blocked_effects_use_four_daemon_execution_rpc_workers() -> None: + runner = ExecutionEffectRunner() + entered = [threading.Event() for _ in range(4)] + release = threading.Event() + try: + for event in entered: + runner.submit_action( + "blocked", + lambda event=event: (event.set(), release.wait())[1], + lambda _: None, + ) + + assert all(event.wait(1) for event in entered) + workers = runner._executor._threads + assert len(workers) == 4 + assert {thread.name for thread in workers} == { + f"execution-rpc_{index}" for index in range(4) + } + assert all(thread.daemon for thread in workers) + finally: + release.set() + runner.shutdown() + + +def test_completion_envelopes_are_immutable_and_correlated() -> None: + runner = ExecutionEffectRunner() + completions: list[object] = [] + complete = threading.Event() + + def enqueue(value: object) -> None: + completions.append(value) + if len(completions) == 4: + complete.set() + + try: + runner.submit_action("action-7", lambda: "done", enqueue) + runner.submit_auxiliary("aux-3", lambda: 42, enqueue) + runner.submit_stop(lambda: None, enqueue) + runner.submit_action("action-8", lambda: None, enqueue) + + assert complete.wait(1) + assert EffectDone("action-7", "done") in completions + assert AuxiliaryDone("aux-3", 42) in completions + assert StopDone(True) in completions + assert EffectDone("action-8", None) in completions + with pytest.raises(FrozenInstanceError): + completions[0].action_id = "changed" # type: ignore[attr-defined] + finally: + runner.shutdown() + + +def test_permanently_blocked_effect_does_not_hold_interpreter_exit() -> None: + code = """ +import threading +from dimos.manipulation.execution_effects import ExecutionEffectRunner + +runner = ExecutionEffectRunner() +runner.submit_action("never", threading.Event().wait, lambda _: None) +""" + process = subprocess.Popen( + [sys.executable, "-c", code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + _, stderr = process.communicate(timeout=2) + except subprocess.TimeoutExpired: + process.kill() + process.communicate() + process.wait() + pytest.fail("permanently blocked effect held interpreter exit") + process.wait() + assert process.returncode == 0, stderr diff --git a/dimos/manipulation/test_execution_gateway.py b/dimos/manipulation/test_execution_gateway.py new file mode 100644 index 0000000000..6d0e57873d --- /dev/null +++ b/dimos/manipulation/test_execution_gateway.py @@ -0,0 +1,59 @@ +"""Black-box tests for coordinator gateway normalization.""" + +from typing import Any + +from dimos.manipulation.execution_gateway import ControlCoordinatorGateway +from dimos.manipulation.execution_models import Outcome +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +class Client: + def __init__(self, value: Any) -> None: + self.value = value + self.calls: list[tuple[str, str, dict[str, Any]]] = [] + + def task_invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + self.calls.append((task, method, kwargs)) + return self.value + + def set_gripper_position(self, hardware_id: str, position: float) -> Any: + self.calls.append((hardware_id, "set_gripper_position", {"position": position})) + return self.value + + def get_gripper_position(self, hardware_id: str) -> Any: + self.calls.append((hardware_id, "get_gripper_position", {})) + return self.value + + +def test_gateway_normalizes_execute_cancel_and_status() -> None: + client = Client(True) + gateway = ControlCoordinatorGateway(client) + assert gateway.execute("task", {"trajectory": "request"}) == Outcome.ACCEPTED + assert gateway.cancel("task") == Outcome.CANCELLED + + client.value = TrajectoryState.EXECUTING + assert gateway.status("task") == Outcome.RUNNING + assert client.calls[0] == ("task", "execute", {"trajectory": "request"}) + + client.value = False + assert gateway.cancel("task") == Outcome.INACTIVE + + +def test_gateway_normalizes_unknown_values_safely() -> None: + client = Client("not-a-state") + gateway = ControlCoordinatorGateway(client) + assert gateway.status("task") == Outcome.UNKNOWN + client.value = object() + assert gateway.execute("task", {}) == Outcome.UNKNOWN + + +def test_gateway_uses_top_level_gripper_rpc_methods() -> None: + client = Client(True) + gateway = ControlCoordinatorGateway(client) + assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED + client.value = 0.6 + assert gateway.get_gripper_position("gripper") == 0.6 + assert client.calls == [ + ("gripper", "set_gripper_position", {"position": 0.4}), + ("gripper", "get_gripper_position", {}), + ] diff --git a/dimos/manipulation/test_execution_models.py b/dimos/manipulation/test_execution_models.py new file mode 100644 index 0000000000..3a834e7404 --- /dev/null +++ b/dimos/manipulation/test_execution_models.py @@ -0,0 +1,27 @@ +"""Public model and snapshot projection tests.""" + +from dataclasses import FrozenInstanceError + +import pytest + +from dimos.manipulation.execution_models import LifecycleState, RuntimeContext, RuntimeSnapshot + + +def test_runtime_snapshot_is_an_immutable_public_projection() -> None: + context = RuntimeContext(state=LifecycleState.READY, diagnostic="ready") + snapshot = RuntimeSnapshot( + context.state, + context.ready_plan, + context.ready_plan_id, + context.planning_token, + context.active, + context.fault, + context.diagnostic, + context.shutdown, + context.shutdown_result, + context.revision, + ) + assert snapshot.state == LifecycleState.READY + assert snapshot.diagnostic == "ready" + with pytest.raises(FrozenInstanceError): + snapshot.diagnostic = "mutated" # type: ignore[misc] diff --git a/dimos/manipulation/test_execution_policy.py b/dimos/manipulation/test_execution_policy.py new file mode 100644 index 0000000000..cc216fcad3 --- /dev/null +++ b/dimos/manipulation/test_execution_policy.py @@ -0,0 +1,116 @@ +import pytest + +from dimos.manipulation.execution_models import ActionMethod, Outcome, TaskActivity +from dimos.manipulation.execution_policy import ( + reconcile_cancel_completion, + reconcile_execute_completion, + reconcile_reset_completion, + reconcile_status_completion, +) + + +@pytest.mark.parametrize( + ("helper", "outcome", "activity"), + [ + (reconcile_execute_completion, Outcome.ACCEPTED, TaskActivity.ACTIVE), + (reconcile_cancel_completion, Outcome.CANCELLED, TaskActivity.CANCELLED), + (reconcile_status_completion, Outcome.COMPLETED, TaskActivity.COMPLETED), + ], +) +def test_completion_policy_normalizes_public_outcomes( + helper: object, outcome: Outcome, activity: TaskActivity +) -> None: + decision = helper(outcome) # type: ignore[operator] + assert decision.activity == activity + assert not decision.fault + + +@pytest.mark.parametrize( + ( + "method", + "outcome", + "activity", + "expected_activity", + "cancel_required", + "request_cancel", + "reset_success", + ), + [ + ( + ActionMethod.STATUS, + Outcome.INACTIVE, + TaskActivity.ACTIVE, + TaskActivity.INACTIVE, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.CANCELLED, + TaskActivity.ACTIVE, + TaskActivity.CANCELLED, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.COMPLETED, + TaskActivity.ACTIVE, + TaskActivity.COMPLETED, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.RUNNING, + TaskActivity.COMPLETED, + TaskActivity.ACTIVE, + True, + True, + None, + ), + ( + ActionMethod.STATUS, + Outcome.ACCEPTED, + TaskActivity.INACTIVE, + TaskActivity.ACTIVE, + True, + True, + None, + ), + ( + ActionMethod.STATUS, + Outcome.UNKNOWN, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + True, + True, + False, + ), + (ActionMethod.CANCEL, Outcome.UNKNOWN, TaskActivity.ACTIVE, None, False, False, False), + (ActionMethod.RESET, Outcome.FAILED, TaskActivity.ACTIVE, None, None, False, False), + ], +) +def test_reset_completion_policy_is_pure_and_classifies_safety( + method: ActionMethod, + outcome: Outcome, + activity: TaskActivity, + expected_activity: TaskActivity | None, + cancel_required: bool | None, + request_cancel: bool, + reset_success: bool | None, +) -> None: + decision = reconcile_reset_completion( + method=method, + outcome=outcome, + activity=activity, + gate_sealed=True, + clock_failed=True, + ) + assert decision.cancel_required == cancel_required + assert decision.request_cancel is request_cancel + assert decision.activity == expected_activity + assert decision.reset_success == reset_success diff --git a/dimos/manipulation/test_execution_runtime.py b/dimos/manipulation/test_execution_runtime.py index 1811063b17..d1e27b8cc1 100644 --- a/dimos/manipulation/test_execution_runtime.py +++ b/dimos/manipulation/test_execution_runtime.py @@ -1,13 +1,14 @@ -"""Focused owner/reconciliation tests.""" +"""Black-box owner, reconciliation, deadline, and shutdown coverage.""" from dataclasses import dataclass, replace from pathlib import Path import threading -import time from typing import Any import pytest +from dimos.manipulation.execution_effects import EffectDone +from dimos.manipulation.execution_models import ActionRecord, Operation, TaskRecord from dimos.manipulation.execution_runtime import ( ActionMethod, ControlCoordinatorGateway, @@ -15,16 +16,13 @@ ExecutionRuntime, ExecutionTopology, LifecycleState, + OperationHandle, Outcome, PreparedPlan, - ResetHandle, - ResetResult, ShutdownState, TaskActivity, TaskEntry, - _Call, - _Event, - _GatewayGateState, + _PendingAction, prepare_generated_plan, ) from dimos.manipulation.planning.groups.models import PlanningGroupDefinition @@ -81,10 +79,12 @@ def __init__( self.entered = {task: threading.Event() for task in executes} self.execute_entered = {task: threading.Event() for task in executes} self.cancel_entered = {task: threading.Event() for task in executes} + self.status_entered = threading.Event() self.release = {task: threading.Event() for task in executes} self._lock = threading.Lock() self.stop_calls = 0 self.calls_at_stop = 0 + self.stop_entered = threading.Event() def _call(self, task: str, method: ActionMethod, outcome: Outcome) -> Outcome: with self._lock: @@ -105,6 +105,9 @@ def cancel(self, task_name: str) -> Outcome: return self._call(task_name, ActionMethod.CANCEL, self.cancels[task_name]) def status(self, task_name: str) -> Outcome: + with self._lock: + self.calls.append((task_name, ActionMethod.STATUS)) + self.status_entered.set() return Outcome.UNKNOWN def reset(self, task_name: str) -> Outcome: @@ -119,6 +122,7 @@ def get_gripper_position(self, hardware_id: str) -> float: return 0.25 def stop(self) -> None: + self.stop_entered.set() with self._lock: self.stop_calls += 1 self.calls_at_stop = len(self.calls) @@ -150,13 +154,6 @@ def plan(*names: str) -> ExecutionPlan: ) -def wait_until(fn: Any, timeout: float = 1.0) -> None: - end = time.monotonic() + timeout - while time.monotonic() < end and not fn(): - time.sleep(0.01) - assert fn() - - def task_is(runtime: ExecutionRuntime, index: int, activity: TaskActivity) -> bool: operation = runtime.snapshot().operation return operation is not None and operation.tasks[index].activity == activity @@ -182,56 +179,19 @@ def test_aggregate_dispatch_pre_registers_and_waits_for_each_acceptance() -> Non runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) try: result = runtime.execute_explicit(plan("a", "b")) - assert result.accepted - - def dispatch_processed() -> bool: - snapshot = runtime.snapshot() - return ( - len([call for call in rpc.calls if call[1] == "execute"]) == 2 - and snapshot.state == LifecycleState.RUNNING - and snapshot.operation is not None - and all(task.activity == TaskActivity.ACTIVE for task in snapshot.operation.tasks) - ) - - wait_until(dispatch_processed) - assert runtime.snapshot().state == LifecycleState.RUNNING + assert result.accepted and result.value is not None + dispatch = runtime.wait_for_dispatch(result.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert len([call for call in rpc.calls if call[1] == "execute"]) == 2 operation = runtime.snapshot().operation assert operation is not None assert len(operation.tasks) == 2 + assert runtime.snapshot().state == LifecycleState.RUNNING + assert all(task.activity == TaskActivity.ACTIVE for task in operation.tasks) finally: runtime.shutdown() -def test_concurrent_execute_admission_is_serial_and_snapshot_coherent() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - barrier = threading.Barrier(3) - results: list[Any] = [] - - def submit() -> None: - barrier.wait() - results.append(runtime.execute_explicit(plan("a"))) - - threads = [threading.Thread(target=submit) for _ in range(2)] - try: - for thread in threads: - thread.start() - barrier.wait() - for thread in threads: - thread.join(1) - assert len(results) == 2 - assert sum(result.accepted for result in results) == 1 - rejected = next(result for result in results if not result.accepted) - assert rejected.value is not None - retained = dict(runtime._context.terminal_results)[rejected.value] - assert retained.outcome == Outcome.REJECTED - assert runtime.snapshot().operation is not None - assert runtime.snapshot().state in (LifecycleState.DISPATCHING, LifecycleState.RUNNING) - finally: - runtime.cancel() - runtime.shutdown() - - def test_ready_replacement_failure_and_consumed_identity() -> None: runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) try: @@ -246,7 +206,6 @@ def test_ready_replacement_failure_and_consumed_identity() -> None: ).accepted assert runtime.snapshot().state == LifecycleState.IDLE assert runtime.snapshot().ready_plan is None - token = runtime.start_planning() assert token is not None assert runtime.complete_planning(token, plan("a")).accepted @@ -261,51 +220,19 @@ def test_ready_replacement_failure_and_consumed_identity() -> None: runtime.shutdown() -def test_active_attempt_admission_rejects_plan_and_execute_but_accepts_cancel() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - try: - planning = runtime.start_planning() - assert planning is not None - assert not runtime.execute_explicit(plan("a")).accepted - assert runtime.cancel_planning().accepted - - started = runtime.execute_explicit(plan("a")) - assert started.accepted - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - assert not runtime.start_planning() - rejected = runtime.execute_explicit(plan("a")) - assert not rejected.accepted and rejected.value is not None - assert dict(runtime._context.terminal_results)[rejected.value].outcome == Outcome.REJECTED - assert runtime.cancel().accepted - assert not runtime.reset().accepted - finally: - gateway.release["a"].set() - runtime.shutdown() - - -def test_rejection_compensates_prior_acceptance_and_unknown_is_fault() -> None: - rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): False, ("a", "cancel"): True}, []) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - runtime.execute_explicit(plan("a", "b")) - wait_until(lambda: any(call[1] == "cancel" for call in rpc.calls)) - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - assert dict(runtime._context.terminal_results).popitem()[1].outcome == Outcome.REJECTED - finally: - runtime.shutdown() - - def test_cancel_during_execute_handles_late_acceptance() -> None: - entered, release = threading.Event(), threading.Event() + entered, release = (threading.Event(), threading.Event()) rpc = FakeRPC({("a", "execute"): True, ("a", "cancel"): True}, [], entered, release) runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) try: - runtime.execute_explicit(plan("a")) + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None assert entered.wait(1) assert runtime.cancel().accepted release.set() - wait_until(lambda: any(call[1] == "cancel" for call in rpc.calls)) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert terminal.accepted and terminal.value is not None + assert any(call[1] == "cancel" for call in rpc.calls) finally: release.set() runtime.shutdown() @@ -333,1479 +260,1751 @@ def test_planning_tokens_and_explicit_execution_are_atomic() -> None: runtime.shutdown() -def test_wait_for_reset_snapshot_is_constructed_outside_context_lock() -> None: - runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) - reset_handle = ResetHandle("reset-attempt") - try: - reset_result = ResetResult(reset_handle, True) - runtime._commit(replace(runtime._context, reset_results=((reset_handle, reset_result),))) - result = runtime.wait_for_reset(reset_handle, timeout=1) - assert result.accepted and result.value == reset_result - finally: - runtime.shutdown() - - def test_deadline_preserves_unresolved_action_and_stop_rpc_is_deferred() -> None: - entered, release = threading.Event(), threading.Event() + clock = AdjustableClock() + entered, release = (threading.Event(), threading.Event()) rpc = FakeRPC({("a", "execute"): True}, [], entered, release) runtime = ExecutionRuntime( - lambda: ControlCoordinatorGateway(rpc), action_timeout=0.01, poll_interval=10 + lambda: ControlCoordinatorGateway(rpc), + monotonic_clock=clock, + action_timeout=1, + poll_interval=100, ) runtime.execute_explicit(plan("a")) assert entered.wait(1) - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) + clock.set(2) + assert not runtime.poll().accepted + assert runtime.snapshot().state == LifecycleState.FAULT assert runtime.snapshot().operation is not None release.set() runtime.shutdown() - assert any(call[1] == "stop" for call in rpc.calls) + threading.Event().wait(0.1) + assert ("", "stop", {}) in rpc.calls -def test_running_inactive_reconciles_peers_before_failed_terminal() -> None: - rpc = FakeRPC( - { - ("a", "execute"): True, - ("b", "execute"): True, - ("a", "get_state"): TrajectoryState.IDLE, - ("b", "get_state"): TrajectoryState.EXECUTING, - ("b", "cancel"): False, - }, - [], +def test_waiters_use_only_the_requested_result_map() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.REJECTED}) + clock = AdjustableClock() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, action_timeout=100, poll_interval=10 ) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) try: - runtime.execute_explicit(plan("a", "b")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime.poll() - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - result = dict(runtime._context.terminal_results).popitem()[1] - assert result.outcome == Outcome.FAILED - assert any(task == "b" and method == "cancel" for task, method, _ in rpc.calls) + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + assert gateway.cancel_entered["a"].wait(1) + dispatch = runtime.wait_for_dispatch(handle, timeout=0.1) + assert ( + dispatch.accepted + and dispatch.value is not None + and (dispatch.value.outcome == Outcome.REJECTED) + ) + terminal = assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(handle, timeout=1) + ) + assert not terminal.accepted and terminal.diagnostic == "timeout" + gateway.release["a"].set() + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert ( + terminal.accepted + and terminal.value is not None + and (terminal.value.outcome == Outcome.REJECTED) + ) finally: + gateway.release["a"].set() runtime.shutdown() -def test_status_unknown_faults_and_cleans_active_peers_without_completion() -> None: - rpc = FakeRPC( - { - ("a", "execute"): True, - ("b", "execute"): True, - ("a", "get_state"): None, - ("b", "cancel"): True, - }, - [], - ) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - runtime.execute_explicit(plan("a", "b")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime.poll() - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - assert not runtime._context.terminal_results - wait_until(lambda: any(task == "b" and method == "cancel" for task, method, _ in rpc.calls)) - finally: - runtime.shutdown() +def test_idle_shutdown_is_idempotent_and_rejects_new_work() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + first = runtime.shutdown(timeout=1) + assert first.accepted and first.value is not None and first.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + assert not getattr(runtime.start_planning(), "accepted", False) + assert not runtime.execute_explicit(plan("a")).accepted + second = runtime.shutdown(timeout=1) + assert second.accepted and second.value == first.value + assert gateway.stop_calls == 1 -def test_dispatch_result_is_stored_once_for_acceptance_and_rejection() -> None: - rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, []) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - handle = runtime.execute_explicit(plan("a", "b")).value - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.ACCEPTED - finally: - runtime.shutdown() +def test_gateway_stop_failure_never_reports_closed_success() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + def fail_stop() -> None: + gateway.stop_calls += 1 + raise RuntimeError("close exploded") -def test_blocked_first_execute_preregisters_batch_and_delays_second() -> None: - entered, release = threading.Event(), threading.Event() - rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, [], entered, release) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - handle = runtime.execute_explicit(plan("a", "b")).value - assert entered.wait(1) - operation = runtime.snapshot().operation - assert operation is not None and len(operation.tasks) == 2 - assert all(task.action is not None for task in operation.tasks[:1]) - assert not any(task == "b" and method == "execute" for task, method, _ in rpc.calls) - assert runtime.snapshot().state == LifecycleState.DISPATCHING - release.set() - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 - finally: - release.set() - runtime.shutdown() + gateway.stop = fail_stop # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None and (not result.value.success) + assert "gateway close failed: close exploded" in result.value.diagnostic + assert result.snapshot is not None + assert result.snapshot.shutdown == ShutdownState.CLOSING + assert result.snapshot.shutdown_result == result.value + assert gateway.stop_calls == 1 + assert runtime.shutdown(timeout=1).accepted is False + assert gateway.stop_calls == 1 -def test_three_task_rejection_skips_unstarted_task_and_rejects_terminally() -> None: - rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): False, ("a", "cancel"): True}, []) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - handle = runtime.execute_explicit(plan("a", "b", "c")).value - assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - assert not any(task == "c" and method == "execute" for task, method, _ in rpc.calls) - assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 - assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.REJECTED - finally: - runtime.shutdown() +def test_invalid_later_clock_sample_does_not_fault_or_kill_owner() -> None: + samples = iter((0.0, 0.0, 0.0, float("nan"))) + invalid_seen = threading.Event() + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + def clock() -> float: + value = next(samples, 0.0) + if value != value: + invalid_seen.set() + return value -def test_cancel_unresolved_dispatch_skips_remaining_tasks() -> None: - entered, release = threading.Event(), threading.Event() - rpc = FakeRPC({("a", "execute"): True, ("a", "cancel"): True}, [], entered, release) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock) try: - handle = runtime.execute_explicit(plan("a", "b", "c")).value - assert handle is not None - assert entered.wait(1) - runtime.cancel() - release.set() - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - assert not any(method == "execute" and task in ("b", "c") for task, method, _ in rpc.calls) - assert len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 - assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.CANCELLED + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert runtime.snapshot().state != LifecycleState.FAULT + assert runtime.snapshot().state == LifecycleState.CANCELLING + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + polled = runtime.poll() + assert not polled.accepted + assert invalid_seen.wait(1) + assert gateway.cancel_entered["a"].wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 1 + gateway.release["a"].set() + assert gateway.stop_entered.wait(1) + assert polled.snapshot is not None + assert polled.snapshot.state != LifecycleState.FAULT finally: - release.set() + gateway.release["a"].set() runtime.shutdown() -def test_deadline_late_accept_and_reject_keep_fault_and_one_dispatch_result() -> None: - for late, expected_calls in ((True, 1), (False, 0)): - entered, release = threading.Event(), threading.Event() - rpc = FakeRPC( - {("a", "execute"): True if late else False, ("a", "cancel"): True}, [], entered, release - ) - runtime = ExecutionRuntime( - lambda rpc=rpc: ControlCoordinatorGateway(rpc), action_timeout=0.01, poll_interval=10 - ) - try: - handle = runtime.execute_explicit(plan("a")).value - assert entered.wait(1) - wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.FAULT) - release.set() - if late: - wait_until(lambda rpc=rpc: any(method == "cancel" for _, method, _ in rpc.calls)) - time.sleep(0.03) - assert runtime.snapshot().state == LifecycleState.FAULT - assert ( - len([item for item in runtime._context.dispatch_results if item[0] == handle]) == 1 - ) - assert len([call for call in rpc.calls if call[1] == "cancel"]) == expected_calls - finally: - release.set() - runtime.shutdown() - +def test_failed_shutdown_drain_rejects_public_commands_after_owner_stops() -> None: + gateway = BlockingGateway({}) -def test_remote_fault_marks_task_and_cancels_active_peer_without_success() -> None: - rpc = FakeRPC( - { - ("a", "execute"): True, - ("b", "execute"): True, - ("a", "get_state"): TrajectoryState.FAULT, - ("b", "cancel"): True, - }, - [], - ) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - runtime.execute_explicit(plan("a", "b")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime.poll() - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None - assert operation.tasks[0].activity.name == "REMOTE_FAULT" - wait_until(lambda: any(task == "b" and method == "cancel" for task, method, _ in rpc.calls)) - assert not runtime._context.terminal_results - finally: - runtime.shutdown() + def fail_stop() -> None: + gateway.stop_calls += 1 + raise RuntimeError("close exploded") - rpc = FakeRPC({("a", "execute"): False}, []) - runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) - try: - handle = runtime.execute_explicit(plan("a")).value - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.REJECTED - finally: - runtime.shutdown() + gateway.stop = fail_stop # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None and not result.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + runtime._owner.join(1) + assert not runtime._owner.is_alive() + + assert not runtime.cancel().accepted + assert not runtime.cancel_if_current(OperationHandle("plan", "operation", "attempt")).accepted + reset = runtime.reset() + assert not reset.accepted + assert reset.diagnostic == "runtime is closing" + assert runtime.shutdown(timeout=1).value == result.value + + +def test_invalid_shutdown_clock_handoff_does_not_deadlock_caller() -> None: + samples = iter((0.0, float("nan"))) + gateway = BlockingGateway({}) + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=lambda: next(samples)) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None + assert result.value.diagnostic == "invalid monotonic clock" + assert runtime.snapshot().shutdown == ShutdownState.CLOSING -def test_cleanup_allocates_distinct_cancel_actions_for_all_active_tasks() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_shutdown_interrupts_early_action_deadline_thread() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b", "c")).value - assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - assert runtime.cancel().accepted - for task in ("a", "b", "c"): - assert gateway.entered[task].wait(1) - operation = runtime.snapshot().operation - assert operation is not None - actions = {task.task_name: task.action for task in operation.tasks} - assert all(actions[name] is not None for name in ("a", "b", "c")) - assert len({actions[name].action_id for name in ("a", "b", "c")}) == 3 # type: ignore[union-attr] - for task in ("a", "b", "c"): - gateway.release[task].set() - wait_until(lambda task=task: gateway.count(task, ActionMethod.CANCEL) == 1) - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - assert dict(runtime._context.terminal_results)[handle].outcome == Outcome.CANCELLED - assert all(gateway.count(task, ActionMethod.CANCEL) == 1 for task in ("a", "b", "c")) + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + gateway.release["a"].set() + assert runtime.shutdown(timeout=1).accepted finally: - runtime.shutdown() + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) -def test_execute_unknown_faults_and_cleans_only_dispatched_execution() -> None: - gateway = BlockingGateway({"a": Outcome.UNKNOWN, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED}) +def test_cancel_if_current_cannot_cancel_replacement_operation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + gateway.status = lambda task_name: Outcome.COMPLETED # type: ignore[method-assign] runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b", "c")).value - assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None and operation.uncertain - assert gateway.entered["a"].wait(1) - assert not any( - task in ("b", "c") and method == ActionMethod.EXECUTE for task, method in gateway.calls - ) - for event in gateway.release.values(): - event.set() - wait_until( - lambda: all(gateway.count(task, ActionMethod.CANCEL) == 1 for task in ("a", "b", "c")) - ) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.UNKNOWN + original = runtime.execute_explicit(plan("a")) + assert original.accepted and original.value is not None + assert runtime.wait_for_dispatch(original.value, timeout=1).accepted + runtime.poll() + assert runtime.wait_for_terminal(original.value, timeout=1).accepted + replacement = runtime.execute_explicit(plan("a")) + assert replacement.accepted and replacement.value is not None + assert runtime.wait_for_dispatch(replacement.value, timeout=1).accepted + stale_cancel = runtime.cancel_if_current(original.value) + assert not stale_cancel.accepted + assert stale_cancel.diagnostic == "operation is no longer current" + active = runtime.snapshot().operation + assert active is not None + assert active.handle == replacement.value + assert gateway.count("a", ActionMethod.CANCEL) == 0 finally: - runtime.shutdown() + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) -def test_cancelled_unresolved_dispatch_retains_cancelled_result_after_late_rejection() -> None: - gateway = BlockingGateway({"a": Outcome.REJECTED}, {"a": Outcome.CANCELLED}, block_execute=True) +def test_shutdown_late_execute_acceptance_is_compensated_before_close() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b")).value - assert handle is not None and gateway.execute_entered["a"].wait(1) - assert runtime.cancel().accepted + runtime.execute_explicit(plan("a")) + assert gateway.execute_entered["a"].wait(1) + result_holder: list[Any] = [] + thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) + thread.start() gateway.release["a"].set() - wait_until(lambda: runtime.snapshot().state == LifecycleState.IDLE) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.CANCELLED - assert gateway.count("a", ActionMethod.CANCEL) == 0 - assert not any( - task == "b" and method == ActionMethod.EXECUTE for task, method in gateway.calls + thread.join(2) + assert ( + len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success ) + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert runtime.snapshot().shutdown == ShutdownState.CLOSED finally: - runtime.shutdown() + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown() -def test_deadline_duplicate_and_late_result_keep_one_original_dispatch_record() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.01, poll_interval=10) - try: - handle = runtime.execute_explicit(plan("a")).value - assert handle is not None - assert gateway.execute_entered["a"].wait(1) - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None and operation.tasks[0].action is not None - task = operation.tasks[0] - action = task.action - assert action is not None - runtime._submit( - lambda: runtime._reduce( - _Event( - "action_deadline", - handle.operation_id, - handle.attempt_id, - task.task_id, - action.action_id, - action.method, - ) - ) - ) - gateway.release["a"].set() - assert gateway.cancel_entered["a"].wait(1) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.UNKNOWN - gateway.release["a"].set() +def _prep_robot( + name: str, + joints: tuple[str, ...], + group: str, + task: str | None, + mapping: dict[str, str] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path(f"/{name}.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), + joint_names=list(joints), + planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], + coordinator_task_name=task, + joint_name_mapping=mapping or {}, + ) + + +def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: + return GeneratedPlan( + group_ids=groups, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width + ), + TrajectoryPoint( + time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width + ), + ], + ), + ) + + +def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + generated = _generated( + ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 + ) + prepared = prepare_generated_plan(generated, topology) + assert prepared.generated_plan is generated + assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] + assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] + assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] + assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 + assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] + + +def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + topology = ExecutionTopology.from_robot_configs((left, right)) + with pytest.raises(ValueError, match="route set"): + ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) + with pytest.raises(ValueError, match="malformed"): + prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) + with pytest.raises(ValueError, match="missing robot joints"): + prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) + with pytest.raises(ValueError, match="duplicate global"): + prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + + +def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + assert isinstance(prepared, PreparedPlan) + assert id(prepared) != id(generated) + partial = PreparedPlan(generated, (), topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + assert not runtime.execute_explicit(partial).accepted finally: - gateway.release["a"].set() runtime.shutdown() -def test_partial_rejection_after_three_accepts_retains_rejected_dispatch_once() -> None: - gateway = BlockingGateway( - {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED, "c": Outcome.ACCEPTED, "d": Outcome.REJECTED} - ) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b", "c", "d")).value - assert handle is not None - wait_until(lambda: gateway.count("d", ActionMethod.EXECUTE) == 1) + legacy = ExecutionPlan(generated, prepared.entries) + assert not runtime.execute_explicit(legacy).accepted + assert runtime.execute_explicit(prepared).accepted + finally: + runtime.shutdown() - def rejected_dispatch_processed() -> bool: - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - return len(records) == 1 and records[0][1].outcome == Outcome.REJECTED - wait_until(rejected_dispatch_processed) - records = [item for item in runtime._context.dispatch_results if item[0] == handle] - assert len(records) == 1 and records[0][1].outcome == Outcome.REJECTED - for event in gateway.release.values(): - event.set() +def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: + left = _prep_robot("left", ("l0",), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + prepared = prepare_generated_plan( + _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology + ) + runtime = ExecutionRuntime( + lambda: BlockingGateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), + poll_interval=10, + ) + try: + result = runtime.execute_explicit(prepared) + assert result.accepted + assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] finally: - for event in gateway.release.values(): - event.set() runtime.shutdown() -def test_late_response_updates_knowledge_without_accepting_stale_action() -> None: - for outcome, expected_activity in ( - (Outcome.ACCEPTED, TaskActivity.CANCELLED), - (Outcome.REJECTED, TaskActivity.INACTIVE), - ): - gateway = BlockingGateway({"a": outcome}, block_execute=True) - runtime = ExecutionRuntime( - lambda gateway=gateway: gateway, action_timeout=0.01, poll_interval=10 +def test_prepare_rejects_additional_selected_robot_joint_column() -> None: + robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + with pytest.raises(ValueError, match="additional joint column"): + prepare_generated_plan( + _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology ) - try: - handle = runtime.execute_explicit(plan("a")).value - assert handle is not None and gateway.execute_entered["a"].wait(1) - wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.FAULT) - gateway.release["a"].set() - if outcome == Outcome.ACCEPTED: - wait_until(lambda gateway=gateway: gateway.count("a", ActionMethod.CANCEL) == 1) - gateway.release["a"].set() - wait_until( - lambda runtime=runtime, expected_activity=expected_activity: task_is( - runtime, 0, expected_activity - ) - ) - operation = runtime.snapshot().operation - assert operation is not None - assert operation.tasks[0].action is None - task_id = operation.tasks[0].task_id - operation_id, attempt_id = handle.operation_id, handle.attempt_id - stale = runtime._submit( - lambda runtime=runtime, - operation_id=operation_id, - attempt_id=attempt_id, - task_id=task_id: runtime._reduce( - _Event( - "action_finished", - operation_id, - attempt_id, - task_id, - "wrong", - ActionMethod.CANCEL, - Outcome.CANCELLED, - ) - ) - ) - assert stale is False - finally: - gateway.release["a"].set() - runtime.shutdown() -def test_remote_fault_requires_reset_and_cancels_peer_without_success() -> None: - gateway = BlockingGateway( - {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, - {"a": Outcome.CANCELLED, "b": Outcome.CANCELLED}, +def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + with pytest.raises(ValueError, match="duplicate group route"): + ExecutionTopology.from_robot_configs( + (robot,), + [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], + ) + with pytest.raises(ValueError, match="mapped local joint"): + ExecutionTopology.from_robot_configs( + (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) + ) + bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + bad_group.planning_groups[0] = PlanningGroupDefinition( + name="manipulator", joint_names=("missing",), base_link="base" ) - gateway.status = lambda task: Outcome.FAILED if task == "a" else Outcome.RUNNING # type: ignore[method-assign] + with pytest.raises(ValueError, match="planning group joints"): + ExecutionTopology.from_robot_configs((bad_group,)) + + +def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: + planning_only = _prep_robot("planner", ("j0",), "manipulator", None) + topology = ExecutionTopology.from_robot_configs((planning_only,)) + assert topology.routes == () + with pytest.raises(ValueError, match="no coordinator task"): + prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) + + +def test_clear_ready_plan_is_a_dedicated_ready_to_idle_transition() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b")).value - assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime.poll() - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None and operation.tasks[0].activity == TaskActivity.REMOTE_FAULT - assert operation.tasks[0].reset_required - assert gateway.cancel_entered["b"].wait(1) - gateway.release["b"].set() - wait_until(lambda: task_is(runtime, 1, TaskActivity.CANCELLED)) - operation = runtime.snapshot().operation - assert operation is not None and all(task.action is None for task in operation.tasks) - assert runtime.snapshot().state == LifecycleState.FAULT - assert not runtime._context.terminal_results + token = runtime.start_planning() + assert token is not None + assert runtime.complete_planning(token, plan("a")).accepted + assert runtime.snapshot().state == LifecycleState.READY + cleared = runtime.clear_ready_plan() + assert cleared.accepted + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.ready_plan is None and snapshot.ready_plan_id is None + assert not runtime.clear_ready_plan().accepted finally: - for event in gateway.release.values(): - event.set() runtime.shutdown() -def test_waiters_use_only_the_requested_result_map() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.REJECTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_runtime_diagnostic_persists_idle_planning_and_execution_errors_then_clears_on_planning() -> ( + None +): + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b")).value - assert handle is not None - assert gateway.cancel_entered["a"].wait(1) - dispatch = runtime.wait_for_dispatch(handle, timeout=0.1) - assert ( - dispatch.accepted - and dispatch.value is not None - and dispatch.value.outcome == Outcome.REJECTED - ) - terminal = runtime.wait_for_terminal(handle, timeout=0.01) - assert not terminal.accepted and terminal.diagnostic == "timeout" - gateway.release["a"].set() - terminal = runtime.wait_for_terminal(handle, timeout=1) - assert ( - terminal.accepted - and terminal.value is not None - and terminal.value.outcome == Outcome.REJECTED - ) + token = runtime.start_planning() + assert token is not None + invalid = runtime.complete_planning(token, ExecutionPlan("generated", ())) + assert not invalid.accepted + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == "invalid prepared plan" + assert runtime.start_planning() is not None + assert runtime.snapshot().diagnostic is None finally: - gateway.release["a"].set() runtime.shutdown() -def test_execute_failure_latches_remote_fault_and_cleans_peer() -> None: - gateway = BlockingGateway({"a": Outcome.FAILED, "b": Outcome.ACCEPTED}) +def test_rejected_execute_retains_contextual_diagnostic_in_snapshot_and_result() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) try: - handle = runtime.execute_explicit(plan("a", "b")).value + handle = runtime.execute_explicit(plan("a")).value assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None - assert operation.tasks[0].activity == TaskActivity.REMOTE_FAULT - assert operation.tasks[0].reset_required - assert gateway.cancel_entered["b"].wait(1) - assert not any( - task == "b" and method == ActionMethod.EXECUTE for task, method in gateway.calls - ) - gateway.release["b"].set() - assert not runtime._context.terminal_results + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.REJECTED + assert "task=a" in result.value.diagnostic + assert "action=" in result.value.diagnostic + assert "outcome=rejected" in result.value.diagnostic + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.diagnostic == result.value.diagnostic finally: - gateway.release["b"].set() runtime.shutdown() -def test_cancel_failure_latches_remote_fault_and_cleans_peer() -> None: - gateway = BlockingGateway( - {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, - {"a": Outcome.FAILED, "b": Outcome.CANCELLED}, - ) +def test_status_terminal_failure_retains_contextual_diagnostic_in_terminal_result() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) + gateway.status = lambda task_name: Outcome.INACTIVE if task_name == "a" else Outcome.RUNNING # type: ignore[assignment,method-assign] + gateway.release["a"].set() + gateway.release["b"].set() runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) try: handle = runtime.execute_explicit(plan("a", "b")).value assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - assert runtime.cancel().accepted - assert gateway.cancel_entered["a"].wait(1) - assert gateway.cancel_entered["b"].wait(1) - gateway.release["a"].set() - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - operation = runtime.snapshot().operation - assert operation is not None - assert operation.tasks[0].activity == TaskActivity.REMOTE_FAULT - assert operation.tasks[0].reset_required - gateway.release["b"].set() - wait_until(lambda: task_is(runtime, 1, TaskActivity.CANCELLED)) - operation = runtime.snapshot().operation - assert operation is not None and all(task.action is None for task in operation.tasks) - assert not runtime._context.terminal_results + assert runtime.wait_for_dispatch(handle, timeout=1).accepted + runtime.poll() + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.FAILED + assert "task=a" in result.value.diagnostic + assert "outcome=inactive" in result.value.diagnostic + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == result.value.diagnostic finally: for event in gateway.release.values(): event.set() runtime.shutdown() -def test_result_retention_is_bounded_and_handles_remain_correlated() -> None: - gateway = BlockingGateway({"a": Outcome.REJECTED}) +def test_auxiliary_gripper_calls_use_private_gateway_and_normalize_results() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - handles = [] try: - for _ in range(20): - result = runtime.execute_explicit(plan("a")) - assert result.accepted and result.value is not None - handle = result.value - handles.append(handle) - dispatch = runtime.wait_for_dispatch(handle, timeout=1) - terminal = runtime.wait_for_terminal(handle, timeout=1) - assert dispatch.accepted and dispatch.value is not None - assert terminal.accepted and terminal.value is not None - assert dispatch.value.handle == handle - assert terminal.value.handle == handle - assert dispatch.value.outcome == Outcome.REJECTED - assert terminal.value.outcome == Outcome.REJECTED - assert len(runtime._context.dispatch_results) == 16 - assert len(runtime._context.terminal_results) == 16 - assert all(handle in dict(runtime._context.dispatch_results) for handle in handles[-16:]) - assert all(handle in dict(runtime._context.terminal_results) for handle in handles[-16:]) - assert handles[0] not in dict(runtime._context.dispatch_results) - assert handles[0] not in dict(runtime._context.terminal_results) + set_result = runtime.set_gripper_position("gripper", 0.75) + get_result = runtime.get_gripper_position("gripper") + assert set_result.accepted and set_result.value == Outcome.ACCEPTED + assert get_result.accepted and get_result.value == 0.25 + assert gateway.count("gripper", ActionMethod.GRIPPER_SET) == 1 + assert gateway.count("gripper", ActionMethod.GRIPPER_GET) == 1 + assert not hasattr(runtime, "gateway") finally: runtime.shutdown() -def test_unknown_fault_reset_fails_and_leaves_fault_latched() -> None: - gateway = BlockingGateway({"a": Outcome.UNKNOWN}) - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.1, poll_interval=10) +def test_control_gateway_uses_top_level_gripper_rpc_methods() -> None: + rpc = FakeRPC( + {("gripper", "set_gripper_position"): True, ("gripper", "get_gripper_position"): 0.6}, [] + ) + gateway = ControlCoordinatorGateway(rpc) + assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED + assert gateway.get_gripper_position("gripper") == 0.6 + assert rpc.calls == [ + ("gripper", "set_gripper_position", {"position": 0.4}), + ("gripper", "get_gripper_position", {}), + ] + + +def test_auxiliary_calls_reject_after_closing_and_inflight_auxiliary_drains_before_close() -> None: + gateway = BlockingAuxGateway() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + results: list[Any] = [] + caller = threading.Thread( + target=lambda: results.append(runtime.set_gripper_position("gripper", 0.5)) + ) + caller.start() + assert gateway.aux_entered.wait(1) + shutdown_results: list[Any] = [] + closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(timeout=1))) + closer.start() + assert closer.is_alive() + gateway.aux_release.set() + caller.join(2) + closer.join(2) + assert results and results[0].accepted + assert shutdown_results and shutdown_results[0].accepted + assert gateway.calls_at_stop == len(gateway.calls) + assert gateway.stop_calls == 1 + assert not runtime.set_gripper_position("gripper", 0.1).accepted + assert not runtime.get_gripper_position("gripper").accepted + + +def test_auxiliary_deadline_waits_for_late_callback_before_shutdown_stop() -> None: + clock = AdjustableClock() + gateway = BlockingAuxGateway() + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + action_timeout=1, + poll_interval=100, + ) + gripper_results: list[Any] = [] + shutdown_results: list[Any] = [] + gripper: threading.Thread | None = None + closer: threading.Thread | None = None + shutdown_started = False try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - assert gateway.cancel_entered["a"].wait(1) - gateway.release["a"].set() - handle_result = runtime.reset() - assert handle_result.accepted and handle_result.value is not None - reset = runtime.wait_for_reset(handle_result.value, timeout=1) - assert reset.accepted and reset.value is not None and not reset.value.success - assert runtime.snapshot().state == LifecycleState.FAULT - time.sleep(0.05) - assert runtime.snapshot().state == LifecycleState.FAULT + gripper = threading.Thread( + target=lambda: gripper_results.append(runtime.set_gripper_position("gripper", 0.5)) + ) + gripper.start() + assert gateway.aux_entered.wait(1) + clock.set(2) + assert not runtime.poll().accepted + gripper.join(1) + assert gripper_results and not gripper_results[0].accepted + + closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(1))) + shutdown_started = True + closer.start() + clock.set(3) + closer.join(1) + assert shutdown_results and not shutdown_results[0].accepted + assert not gateway.stop_entered.is_set() + + gateway.aux_release.set() + assert gateway.stop_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + shutdown_result = runtime.snapshot().shutdown_result + assert shutdown_result is not None and not shutdown_result.success finally: - gateway.release["a"].set() - runtime.shutdown() + gateway.aux_release.set() + if gripper is not None: + gripper.join(2) + if closer is not None: + closer.join(2) + if not shutdown_started: + runtime.shutdown(timeout=1) -def test_remote_fault_reset_proves_inactive_and_resets_remote_task() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) - status_calls: dict[str, int] = {"a": 0, "b": 0} +def test_gateway_factory_is_the_internal_runtime_injection_seam() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) + factories: list[bool] = [] - def status(task: str) -> Outcome: - status_calls[task] += 1 - return Outcome.FAILED if task == "a" and status_calls[task] == 1 else Outcome.INACTIVE + def factory() -> BlockingGateway: + factories.append(True) + return gateway - gateway.status = status # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) + runtime = ExecutionRuntime(factory, poll_interval=10) try: - runtime.execute_explicit(plan("a", "b")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime.poll() - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - gateway.release["b"].set() - wait_until(lambda: gateway.count("b", ActionMethod.CANCEL) == 1) - reset_handle = runtime.reset().value - assert reset_handle is not None - reset = runtime.wait_for_reset(reset_handle, timeout=1) - assert reset.accepted and reset.value is not None and reset.value.success - assert runtime.snapshot().state == LifecycleState.IDLE - assert gateway.count("a", ActionMethod.RESET) == 1 - assert runtime._context.reset_handle is None - assert not runtime._context.reset_proven_inactive - assert not runtime._context.reset_completed_tasks - assert dict(runtime._context.reset_results)[reset_handle] == reset.value + assert factories == [True] + assert runtime.get_gripper_position("gripper").accepted + assert not hasattr(runtime, "client") finally: - for event in gateway.release.values(): - event.set() runtime.shutdown() -def test_remote_reset_failure_keeps_fault_latched() -> None: - for reset_outcome in (Outcome.RUNNING, Outcome.UNKNOWN): - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - gateway.status = lambda task: Outcome.INACTIVE # type: ignore[method-assign] - gateway.reset = lambda task, outcome=reset_outcome: outcome # type: ignore[method-assign] - runtime = ExecutionRuntime( - lambda gateway=gateway: gateway, action_timeout=0.2, poll_interval=10 - ) - try: - runtime.execute_explicit(plan("a")) - wait_until(lambda runtime=runtime: runtime.snapshot().state == LifecycleState.RUNNING) - runtime._submit( - lambda runtime=runtime: runtime._reduce(_Event("fault", diagnostic="remote fault")) - ) - reset_handle = runtime.reset().value - assert reset_handle is not None - reset = runtime.wait_for_reset(reset_handle, timeout=1) - assert reset.accepted and reset.value is not None and not reset.value.success - assert runtime.snapshot().state == LifecycleState.FAULT - finally: - runtime.shutdown() +class AdjustableClock: + def __init__(self) -> None: + self.value = 0.0 + self.lock = threading.Lock() + self._barrier_thread: int | None = None + self._barrier_pending = False + self.sample_entered = threading.Event() + self.sample_release = threading.Event() + + def __call__(self) -> float: + with self.lock: + value = self.value + barrier = self._barrier_pending and self._barrier_thread == threading.get_ident() + self._barrier_pending = False if barrier else self._barrier_pending + if barrier: + self.sample_entered.set() + assert self.sample_release.wait(1) + return value + + def arm_current_thread_sample_barrier(self) -> None: + with self.lock: + self._barrier_thread = threading.get_ident() + self._barrier_pending = True + self.sample_entered.clear() + self.sample_release.clear() + + def set(self, value: float) -> None: + with self.lock: + self.value = value + + +def assert_public_wait_times_out(wait_call: Any) -> Any: + result: list[Any] = [] + ready = threading.Event() + + def observe() -> None: + ready.set() + result.append(wait_call()) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + waiter.join(2) + assert not waiter.is_alive() and len(result) == 1 + return result[0] + + +class ScriptedGateway: + def __init__(self, statuses: dict[str, Outcome] | None = None) -> None: + self.statuses = statuses or {} + self.calls: list[tuple[str, ActionMethod]] = [] + self.lock = threading.Lock() + self.execute_entered = threading.Event() + self.execute_release = threading.Event() + self.block_execute = False + self.cancel_entered = threading.Event() + self.cancel_returned = threading.Event() + self.cancel_entered_by_task: dict[str, threading.Event] = {} + self.cancel_returned_by_task: dict[str, threading.Event] = {} + self.cancel_release = threading.Event() + self.block_cancel = False + self.status_entered = threading.Event() + self.status_release = threading.Event() + self.block_status = False + self.cancel_outcome = Outcome.CANCELLED + self.reset_outcome = Outcome.INACTIVE + self.reset_entered = threading.Event() + self.reset_release = threading.Event() + self.block_reset = False + self.stop_calls = 0 + def _record(self, task: str, method: ActionMethod) -> None: + with self.lock: + self.calls.append((task, method)) -def test_pending_late_execute_action_makes_reset_fail_permanently() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.01, poll_interval=10) - try: - runtime.execute_explicit(plan("a")) - assert gateway.execute_entered["a"].wait(1) - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - reset_handle = runtime.reset().value - assert reset_handle is not None - reset = runtime.wait_for_reset(reset_handle, timeout=1) - assert reset.accepted and reset.value is not None and not reset.value.success - gateway.release["a"].set() - time.sleep(0.05) - assert runtime.snapshot().state == LifecycleState.FAULT - finally: - gateway.release["a"].set() - runtime.shutdown() + def execute(self, task_name: str, request: Any) -> Outcome: + self._record(task_name, ActionMethod.EXECUTE) + self.execute_entered.set() + if self.block_execute: + assert self.execute_release.wait(1) + return Outcome.ACCEPTED + def cancel(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.CANCEL) + self.cancel_entered.set() + self.cancel_entered_by_task.setdefault(task_name, threading.Event()).set() + if self.block_cancel: + assert self.cancel_release.wait(1) + self.cancel_returned.set() + self.cancel_returned_by_task.setdefault(task_name, threading.Event()).set() + return self.cancel_outcome -def test_failed_reset_can_start_a_new_reconciliation_attempt() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - status_calls = 0 + def status(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.STATUS) + self.status_entered.set() + if self.block_status: + assert self.status_release.wait(1) + return self.statuses.get(task_name, Outcome.RUNNING) - def status(_task: str) -> Outcome: - nonlocal status_calls - status_calls += 1 - return Outcome.UNKNOWN if status_calls == 1 else Outcome.INACTIVE + def reset(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.RESET) + self.reset_entered.set() + if self.block_reset: + assert self.reset_release.wait(1) + return self.reset_outcome - gateway.status = status # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) - try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime._submit(lambda: runtime._reduce(_Event("fault", diagnostic="remote fault"))) - gateway.release["a"].set() + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return Outcome.ACCEPTED - first = runtime.reset() - assert first.accepted and first.value is not None - first_result = runtime.wait_for_reset(first.value, timeout=1) - assert first_result.accepted and first_result.value is not None - assert not first_result.value.success - assert runtime.snapshot().state == LifecycleState.FAULT + def get_gripper_position(self, hardware_id: str) -> float: + return 0.0 - second = runtime.reset() - assert second.accepted and second.value is not None - assert second.value != first.value - second_result = runtime.wait_for_reset(second.value, timeout=1) - assert second_result.accepted and second_result.value is not None - assert second_result.value.success - assert runtime.snapshot().state == LifecycleState.IDLE - finally: - gateway.release["a"].set() - runtime.shutdown() + def stop(self) -> None: + self.stop_calls += 1 + self.execute_release.set() + self.cancel_release.set() + self.status_release.set() + def count(self, task: str, method: ActionMethod) -> int: + with self.lock: + return self.calls.count((task, method)) -def test_stale_reset_completion_cannot_clear_fault_for_retry() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - first_reset_entered = threading.Event() - release_first_reset = threading.Event() - second_reset_entered = threading.Event() - release_second_reset = threading.Event() - reset_calls = 0 - - def status(_task: str) -> Outcome: - return Outcome.INACTIVE - - def reset_remote(_task: str) -> Outcome: - nonlocal reset_calls - reset_calls += 1 - if reset_calls == 1: - first_reset_entered.set() - assert release_first_reset.wait(2) - elif reset_calls == 2: - second_reset_entered.set() - assert release_second_reset.wait(2) - return Outcome.INACTIVE - - gateway.status = status # type: ignore[method-assign] - gateway.reset = reset_remote # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) - try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - operation = runtime.snapshot().operation - assert operation is not None - task = operation.tasks[0] - faulted = replace( - operation, - tasks=(replace(task, activity=TaskActivity.INACTIVE, reset_required=True),), - ) - runtime._submit( - lambda: runtime._commit( - replace( - runtime._context, - state=LifecycleState.FAULT, - fault="remote fault", - diagnostic="remote fault", - active=faulted, - ) - ) - ) - first = runtime.reset() - assert first.accepted and first.value is not None - assert first_reset_entered.wait(1) - first_result = runtime.wait_for_reset(first.value, timeout=1) - assert first_result.accepted and first_result.value is not None - assert not first_result.value.success +class SequencedStatusGateway(ScriptedGateway): + def __init__(self, statuses: list[Outcome]) -> None: + super().__init__({"a": Outcome.ACCEPTED}) + self.status_sequence = statuses - second = runtime.reset() - assert second.accepted and second.value is not None and second.value != first.value - release_first_reset.set() - wait_until(lambda: reset_calls == 2) - assert second_reset_entered.wait(1) - assert runtime.snapshot().state == LifecycleState.FAULT - assert runtime._context.reset_handle == second.value + def status(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.STATUS) + self.status_entered.set() + return self.status_sequence.pop(0) - release_second_reset.set() - second_result = runtime.wait_for_reset(second.value, timeout=1) - assert second_result.accepted and second_result.value is not None - assert second_result.value.success - assert runtime.snapshot().state == LifecycleState.IDLE - finally: - release_first_reset.set() - release_second_reset.set() - runtime.shutdown() +def scripted_plan(*names: str) -> ExecutionPlan: + return ExecutionPlan( + "generated", tuple(TaskEntry(name, name, {"trajectory": name}) for name in names) + ) -def test_stale_reset_cancel_cannot_affect_retry_reconciliation() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - first_cancel_entered = threading.Event() - release_first_cancel = threading.Event() - second_cancel_entered = threading.Event() - release_second_cancel = threading.Event() - cancel_calls = 0 - - def cancel(_task: str) -> Outcome: - nonlocal cancel_calls - cancel_calls += 1 - if cancel_calls == 1: - first_cancel_entered.set() - assert release_first_cancel.wait(2) - return Outcome.UNKNOWN - second_cancel_entered.set() - assert release_second_cancel.wait(2) - return Outcome.CANCELLED - - gateway.cancel = cancel # type: ignore[method-assign] - gateway.status = lambda _task: Outcome.INACTIVE # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) - try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - runtime._submit(lambda: runtime._reduce(_Event("fault", diagnostic="remote fault"))) - first = runtime.reset() - assert first.accepted and first.value is not None - assert first_cancel_entered.wait(1) - first_result = runtime.wait_for_reset(first.value, timeout=1) - assert first_result.accepted and first_result.value is not None - assert not first_result.value.success +def start_scripted(runtime: ExecutionRuntime, gateway: ScriptedGateway, task: str = "a") -> Any: + started = runtime.execute_explicit(scripted_plan(task)) + assert started.accepted and started.value is not None + assert gateway.execute_entered.wait(1) + dispatched = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatched.accepted and dispatched.value is not None + assert runtime.snapshot().state == LifecycleState.RUNNING + return started.value - second = runtime.reset() - assert second.accepted and second.value is not None and second.value != first.value - release_first_cancel.set() - wait_until(lambda: cancel_calls == 2) - assert second_cancel_entered.wait(1) - assert runtime.snapshot().state == LifecycleState.FAULT - assert runtime._context.reset_handle == second.value - release_second_cancel.set() - second_result = runtime.wait_for_reset(second.value, timeout=1) - assert second_result.accepted and second_result.value is not None - assert second_result.value.success - assert runtime.snapshot().state == LifecycleState.IDLE - finally: - release_first_cancel.set() - release_second_cancel.set() - runtime.shutdown() +def close_scripted(runtime: ExecutionRuntime, gateway: ScriptedGateway) -> None: + gateway.stop() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) -def test_reset_completion_after_failure_is_retired_before_retry_starts() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - entered = threading.Event() - release = threading.Event() - reset_calls = 0 - - def reset_remote(_task: str) -> Outcome: - nonlocal reset_calls - reset_calls += 1 - if reset_calls == 1: - entered.set() - assert release.wait(2) - return Outcome.INACTIVE - - gateway.reset = reset_remote # type: ignore[method-assign] - gateway.status = lambda _task: Outcome.INACTIVE # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) +@pytest.mark.parametrize("active_status", [Outcome.RUNNING, Outcome.ACCEPTED]) +def test_reset_reconciles_active_status_then_inactive_and_succeeds( + active_status: Outcome, +) -> None: + clock = AdjustableClock() + gateway = SequencedStatusGateway([active_status, Outcome.INACTIVE]) + gateway.cancel_outcome = Outcome.UNKNOWN + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - operation = runtime.snapshot().operation - assert operation is not None - task = operation.tasks[0] - faulted = replace( - operation, - tasks=(replace(task, activity=TaskActivity.INACTIVE, reset_required=True),), - ) - runtime._submit( - lambda: runtime._commit( - replace(runtime._context, state=LifecycleState.FAULT, active=faulted) - ) - ) - first = runtime.reset() - assert first.accepted and first.value is not None and entered.wait(1) - first_result = runtime.wait_for_reset(first.value, timeout=1) - assert first_result.accepted and first_result.value is not None - assert not first_result.value.success - release.set() - wait_until( - lambda: runtime._context.active is not None - and runtime._context.active.tasks[0].action is None - ) + start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + runtime.poll() assert runtime.snapshot().state == LifecycleState.FAULT - second = runtime.reset() - assert second.accepted and second.value is not None and second.value != first.value - wait_until(lambda: reset_calls == 2) - second_result = runtime.wait_for_reset(second.value, timeout=1) - assert second_result.accepted and second_result.value is not None - assert second_result.value.success + gateway.cancel_outcome = Outcome.CANCELLED + gateway.status_entered.clear() + gateway.cancel_returned.clear() + gateway.reset_entered.clear() + reset = runtime.reset() + assert reset.accepted and reset.value is not None + assert gateway.status_entered.wait(1) + assert gateway.count("a", ActionMethod.STATUS) == 1 + gateway.status_entered.clear() + assert gateway.cancel_returned.wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 2 + assert gateway.status_entered.wait(1) + assert gateway.count("a", ActionMethod.STATUS) == 2 + assert gateway.reset_entered.wait(1) + assert gateway.count("a", ActionMethod.RESET) == 1 + result = runtime.wait_for_reset(reset.value, timeout=1) + assert result.accepted and result.value is not None and result.value.success + assert result.snapshot is not None and result.snapshot.state == LifecycleState.IDLE + assert [method for _, method in gateway.calls] == [ + ActionMethod.EXECUTE, + ActionMethod.CANCEL, + ActionMethod.STATUS, + ActionMethod.CANCEL, + ActionMethod.STATUS, + ActionMethod.RESET, + ] finally: - release.set() - runtime.shutdown() + gateway.cancel_release.set() + gateway.reset_release.set() + close_scripted(runtime, gateway) -def test_stale_reset_deadline_does_not_retire_inflight_action() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - entered = threading.Event() - release = threading.Event() - second_entered = threading.Event() - cancel_calls = 0 - - def cancel(_task: str) -> Outcome: - nonlocal cancel_calls - cancel_calls += 1 - if cancel_calls == 2: - second_entered.set() - entered.set() - assert release.wait(2) - return Outcome.CANCELLED - - gateway.cancel = cancel # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.05, poll_interval=10) +def test_queued_poll_uses_injected_clock_and_cancels_due_operation_once() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - operation = runtime.snapshot().operation - assert operation is not None - task = operation.tasks[0] - faulted = replace(operation, tasks=(replace(task, activity=TaskActivity.ACTIVE),)) - runtime._submit( - lambda: runtime._commit( - replace(runtime._context, state=LifecycleState.FAULT, active=faulted) - ) - ) - first = runtime.reset() - assert first.accepted and first.value is not None and entered.wait(1) - old_action = runtime._context.active.tasks[0].action - assert old_action is not None and old_action.reset_id == first.value.reset_id - first_result = runtime.wait_for_reset(first.value, timeout=1) - assert first_result.accepted and first_result.value is not None - assert not first_result.value.success - - second = runtime.reset() - assert second.accepted and second.value is not None - runtime._submit( - lambda: runtime._reduce( - _Event( - "action_deadline", - runtime._context.active.handle.operation_id, - runtime._context.active.handle.attempt_id, - task.task_id, - old_action.action_id, - ActionMethod.CANCEL, - reset_id=first.value.reset_id, - ) - ) - ) - assert runtime._context.active.tasks[0].action is old_action - assert runtime._context.reset_handle == second.value - assert cancel_calls == 1 - assert not second_entered.is_set() - release.set() + start_scripted(runtime, gateway) + clock.set(60) + poll_result: list[Any] = [] + queued = threading.Thread(target=lambda: poll_result.append(runtime.poll())) + queued.start() + queued.join(1) + assert not queued.is_alive() + assert len(poll_result) == 1 + assert poll_result[0].snapshot is not None + assert gateway.count("a", ActionMethod.STATUS) == 0 + assert gateway.count("a", ActionMethod.CANCEL) == 1 finally: - release.set() - runtime.shutdown() - + runtime.shutdown(timeout=0) -def test_repeated_reset_calls_share_one_serialized_result() -> None: - gateway = BlockingGateway({"a": Outcome.UNKNOWN}) - entered, release = threading.Event(), threading.Event() - - def status(task: str) -> Outcome: - entered.set() - assert release.wait(2) - return Outcome.UNKNOWN - gateway.status = status # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, action_timeout=0.2, poll_interval=10) +def test_stale_a_deadline_at_sixty_cannot_cancel_b_started_at_thirty() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.COMPLETED, "b": Outcome.RUNNING}) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.FAULT) - gateway.release["a"].set() - first = runtime.reset() - assert first.accepted and first.value is not None and entered.wait(1) - second = runtime.reset() - assert second.accepted and second.value == first.value - release.set() - result = runtime.wait_for_reset(first.value, timeout=1) - assert result.accepted and result.value is not None and not result.value.success - assert len(runtime._context.reset_results) == 1 + a = start_scripted(runtime, gateway) + runtime.poll() + assert runtime.wait_for_terminal(a, timeout=1).accepted + clock.set(30) + b = start_scripted(runtime, gateway, "b") + gateway.status_entered.clear() + clock.set(61) + runtime.poll() + assert gateway.status_entered.wait(1) + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.handle == b + assert gateway.count("b", ActionMethod.STATUS) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 finally: - release.set() - gateway.release["a"].set() - runtime.shutdown() - - -def test_idle_shutdown_is_idempotent_and_rejects_new_work() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - first = runtime.shutdown(timeout=1) - assert first.accepted and first.value is not None and first.value.success - assert runtime.snapshot().shutdown == ShutdownState.CLOSED - assert not getattr(runtime.start_planning(), "accepted", False) - assert not runtime.execute_explicit(plan("a")).accepted - second = runtime.shutdown(timeout=1) - assert second.accepted and second.value == first.value - assert gateway.stop_calls == 1 + close_scripted(runtime, gateway) -def test_gateway_stop_failure_never_reports_closed_success() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) +def test_cancel_clearance_survives_old_deadline_with_replacement_running() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + a = start_scripted(runtime, gateway) + assert runtime.cancel().accepted + assert gateway.cancel_entered.wait(1) + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.wait_for_terminal(a, timeout=1).accepted + clock.set(30) + start_scripted(runtime, gateway, "b") + clock.set(61) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + finally: + close_scripted(runtime, gateway) - def fail_stop() -> None: - gateway.stop_calls += 1 - raise RuntimeError("close exploded") - gateway.stop = fail_stop # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - result = runtime.shutdown(timeout=1) - assert not result.accepted - assert result.value is not None and not result.value.success - assert "gateway close failed: close exploded" in result.value.diagnostic - assert result.snapshot is not None - assert result.snapshot.shutdown == ShutdownState.CLOSING - assert result.snapshot.shutdown_result == result.value - assert gateway.stop_calls == 1 - assert runtime.shutdown(timeout=1).accepted is False - assert gateway.stop_calls == 1 +def test_reset_rejects_while_fault_cleanup_rpc_is_unresolved() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.FAILED}) + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + start_scripted(runtime, gateway) + runtime.poll() + assert gateway.cancel_entered.wait(1) + snapshot = runtime.snapshot() + assert snapshot.operation is not None + wait = threading.Event() + for _ in range(100): + if runtime.snapshot().state == LifecycleState.FAULT: + break + wait.wait(0.01) + assert runtime.snapshot().state == LifecycleState.FAULT + reset = runtime.reset() + assert not reset.accepted + assert reset.value is None + assert reset.diagnostic == "reset blocked by unresolved coordinator RPC" + assert runtime.snapshot().fault is not None + finally: + gateway.cancel_release.set() + runtime.shutdown(timeout=0) -def test_shutdown_interrupts_early_action_deadline_thread() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) +def test_shutdown_clearance_blocks_cancel_and_returns_truthful_deadline_result() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=100, + poll_interval=100, + ) + result: list[Any] = [] try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - gateway.release["a"].set() - assert runtime.shutdown(timeout=1).accepted - assert runtime._auxiliary == {} + start_scripted(runtime, gateway) + thread = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=100))) + thread.start() + assert gateway.cancel_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + clock.set(61) + gateway.cancel_release.set() + thread.join(1) + assert len(result) == 1 and result[0].value is not None + assert result[0].value.success + assert gateway.count("a", ActionMethod.CANCEL) == 1 finally: - gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown(timeout=1) + gateway.cancel_release.set() + close_scripted(runtime, gateway) -def test_shutdown_interrupts_reset_deadline_thread() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) - handle = ResetHandle("reset-deadline-test") +def test_reset_completes_before_terminal_waiter_observes_retained_diagnostic() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.cancel_outcome = Outcome.UNKNOWN + gateway.statuses["a"] = Outcome.INACTIVE + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime._start_auxiliary( - runtime._reset_deadline, - handle, - time.monotonic() + 10.0, - name="reset-deadline-test", - ) - wait_until(lambda: len(runtime._auxiliary) == 1) - result = runtime.shutdown(timeout=1) - assert result.accepted - assert runtime._auxiliary == {} + handle = start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.snapshot().state == LifecycleState.FAULT + gateway.cancel_outcome = Outcome.CANCELLED + reset = runtime.reset() + assert reset.accepted and reset.value is not None + reset_result = runtime.wait_for_reset(reset.value, timeout=1) + assert reset_result.accepted and reset_result.value is not None + assert reset_result.value.success + waited = runtime.wait_for_terminal(handle, timeout=float("inf")) + assert not waited.accepted + assert waited.diagnostic == "cancellation is uncertain" finally: - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown(timeout=1) + close_scripted(runtime, gateway) -def test_cancel_if_current_cannot_cancel_replacement_operation() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - gateway.status = lambda task_name: Outcome.COMPLETED # type: ignore[method-assign] - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_safe_physical_expiry_returns_failed_terminal_result_and_stable_diagnostic() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - original = runtime.execute_explicit(plan("a")) - assert original.accepted and original.value is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + handle = start_scripted(runtime, gateway) + clock.set(60) runtime.poll() - assert runtime.wait_for_terminal(original.value, timeout=1).accepted + assert gateway.cancel_returned.wait(1) + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.FAILED + assert result.value.diagnostic == "physical operation deadline exceeded" + assert runtime.wait_for_terminal(handle, timeout=1).value == result.value + finally: + close_scripted(runtime, gateway) - replacement = runtime.execute_explicit(plan("a")) - assert replacement.accepted and replacement.value is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - stale_cancel = runtime.cancel_if_current(original.value) - assert not stale_cancel.accepted - assert stale_cancel.diagnostic == "operation is no longer current" - active = runtime.snapshot().operation - assert active is not None - assert active.handle == replacement.value +def test_final_sequential_execute_acceptance_arms_physical_deadline() -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + clock.set(59) + runtime.poll() assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["a"].set() + assert gateway.execute_entered["b"].wait(1) + clock.set(60) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["b"].set() + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + clock.set(120) + runtime.poll() + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) finally: - gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown(timeout=1) + for release in gateway.release.values(): + release.set() + runtime.shutdown() -def test_shutdown_during_pending_cancel_reconciles_and_closes_once() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_blocked_sequential_dispatch_does_not_expire_before_final_acceptance() -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - assert runtime.cancel().accepted - assert gateway.cancel_entered["a"].wait(1) - result_holder: list[Any] = [] - thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) - thread.start() - wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) - assert not runtime.execute_explicit(plan("a")).accepted - assert not getattr(runtime.start_planning(), "accepted", False) + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and gateway.execute_entered["a"].wait(1) + clock.set(600) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 gateway.release["a"].set() - thread.join(2) - assert ( - len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success - ) - assert gateway.stop_calls == 1 + assert gateway.execute_entered["b"].wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["b"].set() + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted # type: ignore[arg-type] + assert gateway.count("a", ActionMethod.CANCEL) == 0 + assert gateway.count("b", ActionMethod.CANCEL) == 0 finally: - gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown() + for release in gateway.release.values(): + release.set() + runtime.shutdown() -def test_shutdown_late_execute_acceptance_is_compensated_before_close() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_finite_terminal_wait_is_observer_only_before_physical_expiry() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime.execute_explicit(plan("a")) - assert gateway.execute_entered["a"].wait(1) - result_holder: list[Any] = [] - thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) - thread.start() - wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) - gateway.release["a"].set() - thread.join(2) - assert ( - len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success - ) - assert gateway.count("a", ActionMethod.CANCEL) == 1 - assert runtime.snapshot().shutdown == ShutdownState.CLOSED + handle = start_scripted(runtime, gateway) + clock.set(59) + observed: list[Any] = [] + ready = threading.Event() + + def observe() -> None: + ready.set() + observed.append(runtime.wait_for_terminal(handle, timeout=1)) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + clock.set(60) + waiter.join(2) + assert not waiter.is_alive() and len(observed) == 1 + waited = observed[0] + assert not waited.accepted and waited.diagnostic == "timeout" + assert gateway.count("a", ActionMethod.CANCEL) == 0 + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) finally: - gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown() + close_scripted(runtime, gateway) -def test_zero_timeout_does_not_publish_closed_while_owner_is_alive() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - entered, release = threading.Event(), threading.Event() - runtime._events.put(_Call(lambda: (entered.set(), release.wait(1)))) +def test_infinite_terminal_wait_returns_after_public_cancellation() -> None: + gateway = ScriptedGateway() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + observed: list[Any] = [] + ready = threading.Event() try: - assert entered.wait(1) - result = runtime.shutdown(timeout=0) - assert not result.accepted and result.value is not None and not result.value.success - assert runtime.snapshot().shutdown == ShutdownState.CLOSING - assert runtime._owner.is_alive() - assert not any( - method in (ActionMethod.EXECUTE, ActionMethod.CANCEL) for _, method in gateway.calls - ) + handle = start_scripted(runtime, gateway) + + def observe() -> None: + ready.set() + observed.append(runtime.wait_for_terminal(handle, timeout=float("inf"))) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + assert runtime.cancel().accepted + waiter.join(1) + assert len(observed) == 1 + assert observed[0].accepted + assert observed[0].value is not None + assert observed[0].value.outcome == Outcome.CANCELLED finally: - release.set() - gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown(timeout=1) + close_scripted(runtime, gateway) -def test_zero_timeout_active_shutdown_never_starts_rpc_after_gateway_stop() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_authoritative_completion_before_physical_expiry_wins() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.COMPLETED}) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime.execute_explicit(plan("a")) - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) - result = runtime.shutdown(timeout=0) - assert not result.accepted and result.value is not None and not result.value.success - time.sleep(0.05) - assert gateway.stop_calls == 1 - assert len(gateway.calls) == gateway.calls_at_stop + handle = start_scripted(runtime, gateway) + clock.set(59) + runtime.poll() + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.COMPLETED + clock.set(60) + runtime.poll() assert gateway.count("a", ActionMethod.CANCEL) == 0 finally: - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown(timeout=1) + close_scripted(runtime, gateway) -def test_gateway_lease_precedes_closure_and_closer_waits_for_inflight_call() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - result_holder: list[tuple[bool, Outcome | None]] = [] +def test_multitask_physical_expiry_cancels_each_task_and_reconciles_terminally() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - thread = threading.Thread( - target=lambda: result_holder.append( - runtime._invoke_gateway_action("action", "a", ActionMethod.EXECUTE, {"x": 1}) - ) - ) - thread.start() - assert gateway.execute_entered["a"].wait(1) - runtime._seal_gateway() - assert gateway.stop_calls == 0 - gateway.release["a"].set() - thread.join(1) - wait_until(lambda: gateway.stop_calls == 1) - assert result_holder == [(True, Outcome.ACCEPTED)] - assert gateway.calls_at_stop == len(gateway.calls) + started = runtime.execute_explicit(scripted_plan("a", "b")) + assert started.accepted and started.value is not None + gateway.cancel_entered_by_task["a"] = threading.Event() + gateway.cancel_entered_by_task["b"] = threading.Event() + gateway.cancel_returned_by_task["a"] = threading.Event() + gateway.cancel_returned_by_task["b"] = threading.Event() + assert gateway.execute_entered.wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + clock.set(60) + runtime.poll() + assert gateway.cancel_returned_by_task["a"].wait(1) + assert gateway.cancel_returned_by_task["b"].wait(1) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert terminal.accepted and terminal.value is not None + assert terminal.value.outcome == Outcome.FAILED finally: - gateway.release["a"].set() - runtime.shutdown(timeout=1) + close_scripted(runtime, gateway) -def test_sealed_gateway_rejects_prepared_action_without_inflight_or_rpc() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +@pytest.mark.parametrize("cancel_outcome", [Outcome.UNKNOWN, Outcome.FAILED]) +def test_physical_expiry_uncertain_cancellation_is_correlated_to_fault( + cancel_outcome: Outcome, +) -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.cancel_outcome = cancel_outcome + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) try: - runtime._seal_gateway() - admitted, outcome = runtime._invoke_gateway_action( - "action", "a", ActionMethod.EXECUTE, {"x": 1} - ) - assert not admitted and outcome is None - assert gateway.calls == [] - assert runtime._rpc_inflight == 0 + handle = start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + observed = runtime.wait_for_terminal(handle, timeout=1) + assert not observed.accepted + assert observed.diagnostic in {"cancellation is uncertain", "cancellation failed"} + assert observed.snapshot is not None + assert observed.snapshot.state == LifecycleState.FAULT finally: - runtime.shutdown(timeout=1) + close_scripted(runtime, gateway) -def test_shutdown_transition_publishes_closing_and_cancel_only_atomically() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - runtime._gateway_condition.acquire() - result_holder: list[Any] = [] - thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) - thread.start() +def test_shutdown_unknown_cancels_each_task_once_and_drains() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, + {"a": Outcome.UNKNOWN, "b": Outcome.UNKNOWN}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + result: list[Any] = [] + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + closer = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=1))) + closer.start() + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + gateway.release["a"].set() + gateway.release["b"].set() + closer.join(2) + assert result and not result[0].accepted + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + assert runtime.snapshot().shutdown_result is not None + assert not runtime.snapshot().shutdown_result.success + finally: + for release in gateway.release.values(): + release.set() + + +def test_shutdown_latches_cancel_only_before_physical_expiry_and_admits_one_cancel() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=100, + poll_interval=100, + ) + result: list[Any] = [] try: - time.sleep(0.02) - assert thread.is_alive() - assert runtime._context.shutdown == ShutdownState.OPEN - assert runtime._gateway_gate == _GatewayGateState.OPEN + start_scripted(runtime, gateway) + thread = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=100))) + thread.start() + assert gateway.cancel_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + assert not runtime.poll().accepted + clock.set(61) + gateway.cancel_release.set() + thread.join(1) + assert result and result[0].value is not None and result[0].value.success + assert gateway.count("a", ActionMethod.CANCEL) == 1 finally: - runtime._gateway_condition.release() - thread.join(2) - assert len(result_holder) == 1 and result_holder[0].accepted - assert runtime.snapshot().shutdown == ShutdownState.CLOSED + gateway.cancel_release.set() + close_scripted(runtime, gateway) -def test_shutdown_timeout_reports_unresolved_action_and_never_schedules_late_effects() -> None: +def test_unresolved_shutdown_deadline_has_no_late_effects() -> None: + clock = AdjustableClock() gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock, poll_interval=100) + shutdown_started = False try: - runtime.execute_explicit(plan("a")) - assert gateway.execute_entered["a"].wait(1) - result = runtime.shutdown(timeout=0.05) - assert not result.accepted and result.value is not None and not result.value.success + started = runtime.execute_explicit(plan("a")) + assert started.accepted and gateway.execute_entered["a"].wait(1) + shutdown_started = True + result = runtime.shutdown(timeout=0) + assert not result.accepted + assert result.value is not None and not result.value.success assert runtime.snapshot().shutdown == ShutdownState.CLOSING cancel_count = gateway.count("a", ActionMethod.CANCEL) gateway.release["a"].set() - time.sleep(0.05) - assert gateway.count("a", ActionMethod.CANCEL) == cancel_count - assert gateway.stop_calls == 1 + assert gateway.stop_entered.wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == cancel_count + 1 finally: gateway.release["a"].set() - if runtime.snapshot().shutdown != ShutdownState.CLOSED: - runtime.shutdown() + if not shutdown_started: + runtime.shutdown(timeout=0) -def _prep_robot( - name: str, - joints: tuple[str, ...], - group: str, - task: str | None, - mapping: dict[str, str] | None = None, -) -> RobotModelConfig: - return RobotModelConfig( - name=name, - model_path=Path(f"/{name}.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), - joint_names=list(joints), - planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], - coordinator_task_name=task, - joint_name_mapping=mapping or {}, +@pytest.mark.parametrize( + ("execute_outcome", "expected_cancel"), + [(Outcome.ACCEPTED, 1), (Outcome.REJECTED, 0)], +) +def test_late_action_deadline_completion_is_correlated( + execute_outcome: Outcome, expected_cancel: int +) -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": execute_outcome}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, action_timeout=1, poll_interval=100 ) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + clock.set(2) + runtime.poll() + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + if expected_cancel: + assert gateway.cancel_entered["a"].wait(1) + gateway.release["a"].set() + assert gateway.count("a", ActionMethod.CANCEL) == expected_cancel + finally: + gateway.release["a"].set() + runtime.shutdown() -def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: - return GeneratedPlan( - group_ids=groups, - trajectory=JointTrajectory( - joint_names=names, - points=[ - TrajectoryPoint( - time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width - ), - TrajectoryPoint( - time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width - ), - ], - ), +def test_unknown_partial_dispatch_compensates_only_admitted_effects() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.UNKNOWN, "c": Outcome.ACCEPTED}, + {"a": Outcome.CANCELLED, "b": Outcome.CANCELLED}, ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b", "c")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + assert not any( + task == "c" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + for release in gateway.release.values(): + release.set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() -def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: - left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) - right = _prep_robot("right", ("r0",), "arm", "right_task") - extra = _prep_robot("extra", ("x0",), "arm", "extra_task") - topology = ExecutionTopology.from_robot_configs((left, right, extra)) - generated = _generated( - ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 +def _install_pending_action( + runtime: ExecutionRuntime, + *, + active_handle: OperationHandle, + pending_handle: OperationHandle, + task_id: str = "task", + active_task_id: str = "task", + current_action_id: str = "new-action", + pending_action_id: str = "old-action", +) -> tuple[ActionRecord, ActionRecord]: + old_action = ActionRecord(pending_action_id, ActionMethod.EXECUTE, 1.0, 1e9) + new_action = ActionRecord(current_action_id, ActionMethod.STATUS, 2.0, 1e9 + 1) + task = TaskRecord( + active_task_id, + "task", + plan("task").entries[0], + TaskActivity.ACTIVE, + new_action, ) - - prepared = prepare_generated_plan(generated, topology) - assert prepared.generated_plan is generated - assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] - assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] - assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] - assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 - assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] - - -def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: - left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") - right = _prep_robot("right", ("r0",), "arm", "right_task") - topology = ExecutionTopology.from_robot_configs((left, right)) - with pytest.raises(ValueError, match="route set"): - ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) - with pytest.raises(ValueError, match="malformed"): - prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) - with pytest.raises(ValueError, match="missing robot joints"): - prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) - with pytest.raises(ValueError, match="duplicate global"): - prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + operation = Operation(active_handle, plan("task"), (task,)) + runtime._commit(state=LifecycleState.RUNNING, active=operation) + runtime._pending[pending_action_id] = _PendingAction( + pending_handle, + "task", + replace(task, task_id=task_id, action=old_action), + ActionMethod.EXECUTE, + 1e9, + ) + runtime._pending[current_action_id] = _PendingAction( + active_handle, + "task", + task, + ActionMethod.STATUS, + 1e9 + 1, + ) + return old_action, new_action -def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: - robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) - topology = ExecutionTopology.from_robot_configs((robot,)) - generated = _generated(["arm/j0"], ("arm/manipulator",), 1) - prepared = prepare_generated_plan(generated, topology) - assert isinstance(prepared, PreparedPlan) - assert prepared is not generated - partial = PreparedPlan(generated, (), topology) - runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) +def test_owner_completion_retires_old_pending_but_preserves_new_action() -> None: + runtime = ExecutionRuntime( + lambda: BlockingGateway({"task": Outcome.RUNNING}), poll_interval=100 + ) try: - assert not runtime.execute_explicit(partial).accepted + handle = OperationHandle("plan", "operation", "attempt") + _, newer = _install_pending_action(runtime, active_handle=handle, pending_handle=handle) + before = runtime.snapshot() + runtime._events.put(EffectDone("old-action", Outcome.ACCEPTED)) + runtime._submit(lambda: None) + after = runtime.snapshot() + assert "old-action" not in runtime._pending + assert "new-action" in runtime._pending + assert after == before + assert after.operation is not None + assert after.operation.tasks[0].action == newer finally: + runtime._pending.clear() + runtime._commit(active=None, state=LifecycleState.IDLE) runtime.shutdown() -def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: - robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") - topology = ExecutionTopology.from_robot_configs((robot,)) - generated = _generated(["arm/j0"], ("arm/manipulator",), 1) - prepared = prepare_generated_plan(generated, topology) - runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) +@pytest.mark.parametrize("kind", ["unknown", "handle", "task"]) +def test_owner_uncorrelatable_completion_preserves_pending_and_snapshot(kind: str) -> None: + runtime = ExecutionRuntime( + lambda: BlockingGateway({"task": Outcome.RUNNING}), poll_interval=100 + ) try: - legacy = ExecutionPlan(generated, prepared.entries) - assert not runtime.execute_explicit(legacy).accepted - assert runtime.execute_explicit(prepared).accepted + handle = OperationHandle("plan", "operation", "attempt") + pending_handle = ( + OperationHandle("other", "operation", "attempt") if kind == "handle" else handle + ) + _install_pending_action( + runtime, + active_handle=handle, + pending_handle=pending_handle, + task_id="missing" if kind == "task" else "task", + ) + before = runtime.snapshot() + action_id = "unknown-action" if kind == "unknown" else "old-action" + runtime._events.put(EffectDone(action_id, Outcome.ACCEPTED)) + runtime._submit(lambda: None) + assert runtime.snapshot() == before + assert set(runtime._pending) == {"old-action", "new-action"} finally: + runtime._pending.clear() + runtime._commit(active=None, state=LifecycleState.IDLE) runtime.shutdown() -def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: - left = _prep_robot("left", ("l0",), "arm", "left_task") - right = _prep_robot("right", ("r0",), "arm", "right_task") - extra = _prep_robot("extra", ("x0",), "arm", "extra_task") - topology = ExecutionTopology.from_robot_configs((left, right, extra)) - prepared = prepare_generated_plan( - _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology - ) - runtime = ExecutionRuntime( - lambda: BlockingGateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), - poll_interval=10, - ) +def test_concurrent_explicit_admission_accepts_one_and_correlates_rejection() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + start = threading.Barrier(3) + results: list[Any] = [] + + def submit() -> None: + start.wait() + results.append(runtime.execute_explicit(plan("a"))) + + threads = [threading.Thread(target=submit) for _ in range(2)] try: - result = runtime.execute_explicit(prepared) - assert result.accepted - assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(1) + assert len(results) == 2 + assert sum(result.accepted for result in results) == 1 + accepted = next(result for result in results if result.accepted) + assert accepted.value is not None + rejected = next(result for result in results if not result.accepted) + assert rejected.value is None + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.handle == accepted.value + gateway.release["a"].set() + assert gateway.count("a", ActionMethod.EXECUTE) == 1 finally: + gateway.release["a"].set() + runtime.cancel() runtime.shutdown() -def test_prepare_rejects_additional_selected_robot_joint_column() -> None: - robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") - topology = ExecutionTopology.from_robot_configs((robot,)) - with pytest.raises(ValueError, match="additional joint column"): - prepare_generated_plan( - _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology - ) +def test_concurrent_ready_execution_consumes_ready_plan_once() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + start = threading.Barrier(3) + results: list[Any] = [] + token = runtime.start_planning() + assert token is not None + completed = runtime.complete_planning(token, plan("a")) + assert completed.accepted and completed.value is not None + plan_id = completed.value + def execute_ready() -> None: + start.wait() + results.append(runtime.execute_ready()) -def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: - robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") - with pytest.raises(ValueError, match="duplicate group route"): - ExecutionTopology.from_robot_configs( - (robot,), - [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], - ) - with pytest.raises(ValueError, match="mapped local joint"): - ExecutionTopology.from_robot_configs( - (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) - ) - bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") - bad_group.planning_groups[0] = PlanningGroupDefinition( - name="manipulator", joint_names=("missing",), base_link="base" - ) - with pytest.raises(ValueError, match="planning group joints"): - ExecutionTopology.from_robot_configs((bad_group,)) + threads = [threading.Thread(target=execute_ready) for _ in range(2)] + try: + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(1) + assert all(not thread.is_alive() for thread in threads) + assert len(results) == 2 + assert sum(result.accepted for result in results) == 1 + assert runtime.snapshot().ready_plan is None + accepted = next(result for result in results if result.accepted) + rejected = next(result for result in results if not result.accepted) + assert accepted.value is not None + assert accepted.value.plan_id == plan_id + assert not rejected.accepted + assert rejected.value is None + assert rejected.diagnostic == "no ready plan" + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + finally: + gateway.release["a"].set() + runtime.cancel() + runtime.shutdown() -def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: - planning_only = _prep_robot("planner", ("j0",), "manipulator", None) - topology = ExecutionTopology.from_robot_configs((planning_only,)) - assert topology.routes == () - with pytest.raises(ValueError, match="no coordinator task"): - prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) +def test_active_operation_rejects_planning_and_explicit_execution_but_accepts_cancel() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + original = runtime.snapshot() + assert original.state == LifecycleState.RUNNING + assert not runtime.start_planning() + rejected = runtime.execute_explicit(plan("a")) + assert not rejected.accepted + assert rejected.value is None + assert runtime.snapshot().operation == original.operation + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert runtime.snapshot().operation == original.operation + assert runtime.cancel().accepted + gateway.release["a"].set() + finally: + gateway.release["a"].set() + runtime.shutdown() -def test_clear_ready_plan_is_a_dedicated_ready_to_idle_transition() -> None: - gateway = BlockingGateway({"a": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_sequential_rejection_skips_remaining_task_and_compensates_prior_acceptance() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.REJECTED, "c": Outcome.ACCEPTED}, + {"a": Outcome.CANCELLED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) try: - token = runtime.start_planning() - assert token is not None - assert runtime.complete_planning(token, plan("a")).accepted - assert runtime.snapshot().state == LifecycleState.READY - cleared = runtime.clear_ready_plan() - assert cleared.accepted - snapshot = runtime.snapshot() - assert snapshot.state == LifecycleState.IDLE - assert snapshot.ready_plan is None and snapshot.ready_plan_id is None - assert not runtime.clear_ready_plan().accepted + started = runtime.execute_explicit(plan("a", "b", "c")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) + assert gateway.calls[:3] == [ + ("a", ActionMethod.EXECUTE), + ("b", ActionMethod.EXECUTE), + ("a", ActionMethod.CANCEL), + ] + assert not any( + task == "c" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.outcome == Outcome.REJECTED + assert terminal.value.outcome == Outcome.REJECTED + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 1 + assert gateway.count("c", ActionMethod.EXECUTE) == 0 + assert gateway.count("c", ActionMethod.CANCEL) == 0 finally: + for release in gateway.release.values(): + release.set() runtime.shutdown() -def test_runtime_diagnostic_persists_idle_planning_and_execution_errors_then_clears_on_planning() -> ( - None -): - runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) +def test_cancel_during_unresolved_dispatch_retains_cancelled_result_after_late_acceptance() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) try: - token = runtime.start_planning() - assert token is not None - invalid = runtime.complete_planning(token, ExecutionPlan("generated", ())) - assert not invalid.accepted - assert runtime.snapshot().state == LifecycleState.IDLE - assert runtime.snapshot().diagnostic == "invalid prepared plan" - assert runtime.start_planning() is not None - assert runtime.snapshot().diagnostic is None + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.cancel().accepted + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.outcome == Outcome.CANCELLED + assert terminal.value.outcome == Outcome.CANCELLED + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.calls[:2] == [ + ("a", ActionMethod.EXECUTE), + ("a", ActionMethod.CANCEL), + ] + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 0 + assert gateway.count("b", ActionMethod.CANCEL) == 0 finally: + for release in gateway.release.values(): + release.set() runtime.shutdown() -def test_rejected_execute_retains_contextual_diagnostic_in_snapshot_and_result() -> None: - gateway = BlockingGateway({"a": Outcome.REJECTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_unknown_execute_enters_fault_and_compensates_without_terminal_success() -> None: + gateway = BlockingGateway( + {"a": Outcome.UNKNOWN, "b": Outcome.REJECTED}, + {"a": Outcome.UNKNOWN, "b": Outcome.FAILED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) try: - handle = runtime.execute_explicit(plan("a")).value - assert handle is not None - result = runtime.wait_for_terminal(handle, timeout=1) - assert result.accepted and result.value is not None - assert result.value.outcome == Outcome.REJECTED - assert "task=a" in result.value.diagnostic - assert "action=" in result.value.diagnostic - assert "outcome=rejected" in result.value.diagnostic + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) snapshot = runtime.snapshot() - assert snapshot.state == LifecycleState.IDLE - assert snapshot.diagnostic == result.value.diagnostic + assert snapshot.state == LifecycleState.FAULT + assert snapshot.operation is not None and snapshot.operation.uncertain + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 0 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + assert dispatch.value.diagnostic + assert "task=a" in dispatch.value.diagnostic + assert "outcome=unknown" in dispatch.value.diagnostic + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert not terminal.accepted + assert terminal.diagnostic == dispatch.value.diagnostic finally: + gateway.release["a"].set() runtime.shutdown() -def test_status_terminal_failure_retains_contextual_diagnostic_in_terminal_result() -> None: +def test_unknown_status_enters_fault_without_successful_terminal_result() -> None: gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) - gateway.status = lambda task: Outcome.INACTIVE if task == "a" else Outcome.RUNNING # type: ignore[method-assign] - gateway.release["a"].set() - gateway.release["b"].set() - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) try: - handle = runtime.execute_explicit(plan("a", "b")).value - assert handle is not None - wait_until(lambda: runtime.snapshot().state == LifecycleState.RUNNING) + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + handle = started.value + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted runtime.poll() - result = runtime.wait_for_terminal(handle, timeout=1) - assert result.accepted and result.value is not None - assert result.value.outcome == Outcome.FAILED - assert "task=a" in result.value.diagnostic - assert "outcome=inactive" in result.value.diagnostic - assert runtime.snapshot().state == LifecycleState.IDLE - assert runtime.snapshot().diagnostic == result.value.diagnostic + assert gateway.status_entered.wait(1) + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + gateway.release["a"].set() + gateway.release["b"].set() + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert not terminal.accepted + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert terminal.diagnostic == runtime.snapshot().diagnostic + assert terminal.diagnostic + assert "task=a" in terminal.diagnostic + assert "outcome=unknown" in terminal.diagnostic + assert "status is unsafe" in terminal.diagnostic + assert runtime.snapshot().state == LifecycleState.FAULT + assert gateway.count("a", ActionMethod.STATUS) == 1 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 finally: - for event in gateway.release.values(): - event.set() + for release in gateway.release.values(): + release.set() runtime.shutdown() -def test_auxiliary_gripper_calls_use_private_gateway_and_normalize_results() -> None: - gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) +def test_failed_cancel_latches_fault_and_preserves_reset_required_operation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, {"a": Outcome.FAILED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) try: - set_result = runtime.set_gripper_position("gripper", 0.75) - get_result = runtime.get_gripper_position("gripper") - assert set_result.accepted and set_result.value == Outcome.ACCEPTED - assert get_result.accepted and get_result.value == 0.25 - assert gateway.count("gripper", ActionMethod.GRIPPER_SET) == 1 - assert gateway.count("gripper", ActionMethod.GRIPPER_GET) == 1 - assert not hasattr(runtime, "gateway") + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + assert runtime.cancel().accepted + assert gateway.cancel_entered["a"].wait(1) + gateway.release["a"].set() + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert not terminal.accepted + assert terminal.diagnostic == "cancellation failed" + assert runtime.snapshot().state == LifecycleState.FAULT + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.tasks[0].reset_required + assert snapshot.diagnostic == "cancellation failed" + assert gateway.count("a", ActionMethod.CANCEL) == 1 finally: + gateway.release["a"].set() runtime.shutdown() -def test_control_gateway_uses_top_level_gripper_rpc_methods() -> None: - rpc = FakeRPC( - {("gripper", "set_gripper_position"): True, ("gripper", "get_gripper_position"): 0.6}, [] - ) - gateway = ControlCoordinatorGateway(rpc) - assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED - assert gateway.get_gripper_position("gripper") == 0.6 - assert rpc.calls == [ - ("gripper", "set_gripper_position", {"position": 0.4}), - ("gripper", "get_gripper_position", {}), - ] +def test_reset_calls_share_one_public_handle_until_first_result_is_ready() -> None: + gateway = ScriptedGateway({"a": Outcome.FAILED}) + gateway.block_reset = True + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + start_scripted(runtime, gateway) + runtime.poll() + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert not runtime.wait_for_terminal(snapshot.operation.handle, timeout=1).accepted + gateway.statuses["a"] = Outcome.INACTIVE + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.snapshot().state == LifecycleState.FAULT + first = runtime.reset() + assert first.accepted and first.value is not None + assert gateway.reset_entered.wait(1) + second = runtime.reset() + assert second.accepted and second.value == first.value + gateway.reset_release.set() + result = runtime.wait_for_reset(first.value, timeout=1) + assert result.accepted and result.value is not None and result.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + assert gateway.count("a", ActionMethod.RESET) == 1 + finally: + gateway.reset_release.set() + close_scripted(runtime, gateway) -def test_auxiliary_calls_reject_after_closing_and_inflight_auxiliary_drains_before_close() -> None: - gateway = BlockingAuxGateway() - runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) - results: list[Any] = [] - caller = threading.Thread( - target=lambda: results.append(runtime.set_gripper_position("gripper", 0.5)) - ) - caller.start() - assert gateway.aux_entered.wait(1) - shutdown_results: list[Any] = [] - closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(timeout=1))) - closer.start() - wait_until(lambda: runtime.snapshot().shutdown == ShutdownState.CLOSING) - gateway.aux_release.set() - caller.join(2) - closer.join(2) - assert results and results[0].accepted - assert shutdown_results and shutdown_results[0].accepted - assert gateway.calls_at_stop == len(gateway.calls) - assert gateway.stop_calls == 1 - assert not runtime.set_gripper_position("gripper", 0.1).accepted - assert not runtime.get_gripper_position("gripper").accepted +def test_unsafe_reset_remains_faulted_without_an_in_process_retry() -> None: + gateway = ScriptedGateway() + gateway.block_cancel = True + status_calls = 0 + def status(task_name: str) -> Outcome: + nonlocal status_calls + status_calls += 1 + return Outcome.FAILED if status_calls == 1 else Outcome.UNKNOWN -def test_gateway_factory_is_the_internal_runtime_injection_seam() -> None: - gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) - factories: list[bool] = [] + gateway.status = status # type: ignore[assignment,method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + start_scripted(runtime, gateway) + runtime.poll() + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert not runtime.wait_for_terminal(snapshot.operation.handle, timeout=1).accepted + assert gateway.cancel_entered.wait(1) + reset = runtime.reset() + assert not reset.accepted and reset.value is None + assert reset.diagnostic == "reset blocked by unresolved coordinator RPC" + assert runtime.snapshot().state == LifecycleState.FAULT + assert runtime.snapshot().diagnostic + assert gateway.count("a", ActionMethod.RESET) == 0 + finally: + gateway.cancel_release.set() + close_scripted(runtime, gateway) - def factory() -> BlockingGateway: - factories.append(True) - return gateway - runtime = ExecutionRuntime(factory, poll_interval=10) +def test_public_results_retain_current_and_recent_handles_and_age_out_old_handles() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) + clock = AdjustableClock() + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock, poll_interval=100) + handles: list[Any] = [] try: - assert factories == [True] - assert runtime.get_gripper_position("gripper").accepted - assert not hasattr(runtime, "client") + for _ in range(20): + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + handles.append(started.value) + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.handle == started.value + assert terminal.value.handle == started.value + for handle in handles[-2:]: + dispatch = runtime.wait_for_dispatch(handle, timeout=1) + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert dispatch.accepted and terminal.accepted + assert dispatch.value is not None and dispatch.value.outcome == Outcome.REJECTED + assert terminal.value is not None and terminal.value.outcome == Outcome.REJECTED + assert dispatch.value.handle == handle + assert terminal.value.handle == handle + aged_out = handles[0] + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(aged_out, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(aged_out, timeout=1) + ).diagnostic + == "timeout" + ) + wrong = replace(handles[-1], attempt_id="missing-attempt") + unknown = replace(handles[-1], operation_id="missing-operation") + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(wrong, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(wrong, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(unknown, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(unknown, timeout=1) + ).diagnostic + == "timeout" + ) finally: runtime.shutdown() diff --git a/dimos/manipulation/test_execution_topology.py b/dimos/manipulation/test_execution_topology.py new file mode 100644 index 0000000000..c6478633f8 --- /dev/null +++ b/dimos/manipulation/test_execution_topology.py @@ -0,0 +1,199 @@ +"""Topology and generated-plan materialization tests.""" + +from pathlib import Path + +import pytest + +from dimos.manipulation.execution_models import Outcome +from dimos.manipulation.execution_runtime import ExecutionRuntime, ExecutionTopology +from dimos.manipulation.execution_topology import ( + ExecutionPlan, + PreparedPlan, + prepare_generated_plan, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +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.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +class Gateway: + def __init__(self, outcomes: dict[str, Outcome]) -> None: + self.outcomes = outcomes + + def execute(self, task_name: str, request: object) -> Outcome: + return self.outcomes.get(task_name, Outcome.ACCEPTED) + + def cancel(self, task_name: str) -> Outcome: + return Outcome.CANCELLED + + def status(self, task_name: str) -> Outcome: + return Outcome.INACTIVE + + def reset(self, task_name: str) -> Outcome: + return Outcome.INACTIVE + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return Outcome.ACCEPTED + + def get_gripper_position(self, hardware_id: str) -> float: + return 0.0 + + def stop(self) -> None: + return None + + +def _prep_robot( + name: str, + joints: tuple[str, ...], + group: str, + task: str | None, + mapping: dict[str, str] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path(f"/{name}.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), + joint_names=list(joints), + planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], + coordinator_task_name=task, + joint_name_mapping=mapping or {}, + ) + + +def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: + return GeneratedPlan( + group_ids=groups, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width + ), + TrajectoryPoint( + time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width + ), + ], + ), + ) + + +def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + generated = _generated( + ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 + ) + + prepared = prepare_generated_plan(generated, topology) + assert prepared.generated_plan is generated + assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] + assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] + assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] + assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 + assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] + + +def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + topology = ExecutionTopology.from_robot_configs((left, right)) + with pytest.raises(ValueError, match="route set"): + ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) + with pytest.raises(ValueError, match="malformed"): + prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) + with pytest.raises(ValueError, match="missing robot joints"): + prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) + with pytest.raises(ValueError, match="duplicate global"): + prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + + +def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: + left = _prep_robot("left", ("l0",), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + prepared = prepare_generated_plan( + _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology + ) + runtime = ExecutionRuntime( + lambda: Gateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), + poll_interval=10, + ) + try: + result = runtime.execute_explicit(prepared) + assert result.accepted + assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] + finally: + runtime.shutdown() + + +def test_prepare_rejects_additional_selected_robot_joint_column() -> None: + robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + with pytest.raises(ValueError, match="additional joint column"): + prepare_generated_plan( + _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology + ) + + +def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + with pytest.raises(ValueError, match="duplicate group route"): + ExecutionTopology.from_robot_configs( + (robot,), + [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], + ) + with pytest.raises(ValueError, match="mapped local joint"): + ExecutionTopology.from_robot_configs( + (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) + ) + bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + bad_group.planning_groups[0] = PlanningGroupDefinition( + name="manipulator", joint_names=("missing",), base_link="base" + ) + with pytest.raises(ValueError, match="planning group joints"): + ExecutionTopology.from_robot_configs((bad_group,)) + + +def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: + planning_only = _prep_robot("planner", ("j0",), "manipulator", None) + topology = ExecutionTopology.from_robot_configs((planning_only,)) + assert topology.routes == () + with pytest.raises(ValueError, match="no coordinator task"): + prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) + + +def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + assert isinstance(prepared, PreparedPlan) + assert prepared is not generated + partial = PreparedPlan(generated, (), topology) + runtime = ExecutionRuntime(lambda: Gateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + assert not runtime.execute_explicit(partial).accepted + finally: + runtime.shutdown() + + +def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + runtime = ExecutionRuntime(lambda: Gateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + legacy = ExecutionPlan(generated, prepared.entries) + assert not runtime.execute_explicit(legacy).accepted + assert runtime.execute_explicit(prepared).accepted + finally: + runtime.shutdown() diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index b1f49344dd..bb1a24abd5 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -22,9 +22,13 @@ from __future__ import annotations import importlib.util +import threading +from types import SimpleNamespace +from typing import Any import pytest +from dimos.manipulation import manipulation_module as module_impl from dimos.manipulation._test_manipulation_helpers import FakeCoordinatorGateway from dimos.manipulation.execution_runtime import ( ExecutionRuntime, @@ -135,6 +139,157 @@ def module(xarm7_config): mod.stop() +def test_execute_explicit_plan_cannot_dispatch_replaced_ready_plan(monkeypatch) -> None: + """A READY replacement between observation and submission cannot switch plans.""" + + plan_a = object() + plan_b = object() + prepared_a = object() + ready_b = PreparedPlan(generated_plan=plan_b, entries=()) # type: ignore[arg-type] + ready_a = PreparedPlan(generated_plan=plan_a, entries=()) # type: ignore[arg-type] + dispatched: list[object] = [] + ready_dispatches: list[object] = [] + + class FakeRuntime: + def snapshot(self): + # Model the replacement at precisely the old snapshot/command seam. + self.current_ready = ready_b + return SimpleNamespace(ready_plan=ready_a) + + def execute_ready(self): + ready_dispatches.append(self.current_ready.generated_plan) + return SimpleNamespace(accepted=True, value=object()) + + def execute_explicit(self, prepared): + dispatched.append(prepared) + return SimpleNamespace(accepted=True, value=object()) + + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = object() + monkeypatch.setattr(module_impl, "prepare_generated_plan", lambda plan, topology: prepared_a) + + assert module._submit_execution(plan_a) is not None + assert dispatched == [prepared_a] + assert ready_dispatches == [] + + +def test_execute_waits_for_its_correlated_dispatch_result() -> None: + released = threading.Event() + entered = threading.Event() + expected_handle = object() + waits: list[object] = [] + + class FakeRuntime: + def execute_ready(self): + return SimpleNamespace(accepted=True, value=expected_handle) + + def wait_for_dispatch(self, handle, *, timeout): + waits.append(handle) + entered.set() + assert released.wait(1) + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.ACCEPTED), + ) + + module = object.__new__(ManipulationModule) + module._execution_runtime = FakeRuntime() + module._execution_topology = object() + result: list[bool] = [] + caller = threading.Thread(target=lambda: result.append(module.execute())) + caller.start() + assert entered.wait(1) + assert result == [] + released.set() + caller.join(1) + assert result == [True] + assert waits == [expected_handle] + + +def test_preview_waits_for_terminal_before_reporting_completion(monkeypatch) -> None: + released = threading.Event() + entered = threading.Event() + expected_handle = object() + preview_calls: list[tuple[object, object]] = [] + + class FakeRuntime: + def execute_explicit(self, prepared): + return SimpleNamespace(accepted=True, value=expected_handle) + + def wait_for_dispatch(self, handle, *, timeout): + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.ACCEPTED), + ) + + def wait_for_terminal(self, handle, *, timeout): + entered.set() + assert released.wait(1) + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.COMPLETED, diagnostic=""), + ) + + plan = object() + ready = PreparedPlan(generated_plan=plan, entries=()) # type: ignore[arg-type] + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = object() + module.preview_path = lambda duration=None, robot_name=None, target_fps=30: ( + preview_calls.append((duration, robot_name)) or True + ) + monkeypatch.setattr(module_impl, "prepare_generated_plan", lambda value, topology: value) + + # The preview helper obtains the READY generated plan only to bind it through + # the runtime's explicit-plan path; terminal publication is the completion gate. + runtime.snapshot = lambda: SimpleNamespace(ready_plan=ready) + result: list[Any] = [] + caller = threading.Thread(target=lambda: result.append(module._preview_execute_wait())) + caller.start() + assert entered.wait(1) + assert result == [] + assert preview_calls == [(0.5, None)] + released.set() + caller.join(1) + assert len(result) == 1 and result[0].is_success() + + +def test_execution_projection_helpers_use_runtime_snapshot_only() -> None: + snapshot = SimpleNamespace( + state=LifecycleState.READY, + diagnostic="runtime diagnostic", + ready_plan_id="ready-1", + ready_plan=PreparedPlan(generated_plan=SimpleNamespace(path=[]), entries=()), + operation=None, + ) + + class FakeRuntime: + def __init__(self) -> None: + self.snapshot_calls = 0 + + def snapshot(self): + self.snapshot_calls += 1 + return snapshot + + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = None + + assert module.get_state() == LifecycleState.READY.name + assert module.get_error() == "runtime diagnostic" + assert module.has_planned_path() is False + projection = module.get_execution_snapshot() + assert projection.ready_plan_id == "ready-1" + assert projection.has_ready_plan is True + assert runtime.snapshot_calls == 4 + assert not hasattr(module, "_execution_state") + assert not hasattr(module, "_execution_error") + + @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") class TestManipulationModuleIntegration: @@ -292,9 +447,6 @@ def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): assert result is False assert module.get_state() == LifecycleState.IDLE.name assert "rejected" in module.get_error().lower() - terminal_results = dict(module._execution_runtime._context.terminal_results) - assert terminal_results - assert next(reversed(terminal_results.values())).outcome == Outcome.REJECTED def test_state_transitions_during_execution(self, module, joint_state_zeros): """Test state transitions during plan and execute.""" diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py index a056351c5f..3bf29c883b 100644 --- a/dimos/manipulation/test_plan_execution_reservation.py +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -9,18 +9,26 @@ from pathlib import Path from unittest.mock import MagicMock +from pydantic import ValidationError import pytest from dimos.manipulation._test_manipulation_helpers import ( FakeCoordinatorGateway, close_test_runtimes, + install_ready_plan, install_runtime, make_module, ) -from dimos.manipulation.execution_runtime import LifecycleState, Outcome, prepare_generated_plan +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + LifecycleState, + Outcome, + prepare_generated_plan, +) from dimos.manipulation.manipulation_module import ( ManipulationExecutionSnapshot, ManipulationModule, + ManipulationModuleConfig, ) from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry @@ -34,6 +42,112 @@ from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +def test_physical_operation_timeout_config_defaults_to_sixty_and_rejects_invalid_values() -> None: + assert ManipulationModuleConfig().physical_operation_timeout == 60.0 + for value in (0.0, -1.0, float("inf"), float("nan")): + with pytest.raises(ValidationError): + ManipulationModuleConfig(physical_operation_timeout=value) + + +def test_module_runtime_factory_receives_distinct_action_and_physical_timeouts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = ManipulationModule( + robots=[_config()], + planning_timeout=3.0, + physical_operation_timeout=17.0, + visualization={"backend": "none"}, + ) + module.coordinator_joint_state = None + captured: dict[str, float] = {} + runtimes: list[ExecutionRuntime] = [] + + class WorldMonitor: + visualization = None + + def add_robot(self, config: RobotModelConfig) -> str: + del config + return "robot-id" + + def finalize(self) -> None: ... + + def start_state_monitor(self, robot_id: str) -> None: + del robot_id + + def set_visualization(self, visualization: object) -> None: + del visualization + + def initialize_visualization(self, operator: object) -> None: + del operator + + def stop_all_monitors(self) -> None: ... + + world_monitor = WorldMonitor() + specs = type( + "PlanningSpecs", + (), + {"world_monitor": world_monitor, "planner": object(), "kinematics": object()}, + )() + + def runtime_factory(_gateway_factory, **kwargs): + captured.update(kwargs) + runtime = ExecutionRuntime(lambda: FakeCoordinatorGateway(), **kwargs) + runtimes.append(runtime) + return runtime + + monkeypatch.setattr("dimos.manipulation.manipulation_module.create_world", lambda **_: object()) + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.create_planning_specs", + lambda **_: specs, + ) + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.create_manipulation_visualization", + lambda *_, **__: None, + ) + monkeypatch.setattr(ManipulationModule, "_runtime_factory", staticmethod(runtime_factory)) + + try: + module.start() + assert captured["action_timeout"] == 3.0 + assert captured["physical_operation_timeout"] == 17.0 + finally: + module.stop() + + +def test_public_move_to_joints_reports_timeout_fault_without_extra_cancel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + gateway = FakeCoordinatorGateway( + status_outcome=Outcome.RUNNING, + status_outcomes=[Outcome.RUNNING], + cancel_outcome=Outcome.UNKNOWN, + ) + module, _ = _module(physical_operation_timeout=0.02, gateway=gateway) + plan = _plan() + + def plan_to_joints(_goal, _robot_name=None) -> bool: + install_ready_plan(module, plan) + return True + + monkeypatch.setattr(module, "plan_to_joints", plan_to_joints) + monkeypatch.setattr(module, "preview_path", lambda *args, **kwargs: True) + + result = module.move_to_joints("1.0") + + assert not result.is_success() + assert result.error_code == "EXECUTION_FAILED" + assert "uncertain" in result.message.lower() + assert module.get_state() == LifecycleState.FAULT.name + assert module.get_error() == result.message + snapshot = module.get_execution_snapshot() + assert snapshot.state == LifecycleState.FAULT.name + assert snapshot.diagnostic == result.message + assert gateway.cancel_calls == ["traj_arm"] + + module.get_state() + assert gateway.cancel_calls == ["traj_arm"] + + def _config(task_name: str | None = "traj_arm") -> RobotModelConfig: return RobotModelConfig( name="arm", @@ -71,14 +185,16 @@ def _plan(group_id: str = "arm/manipulator") -> GeneratedPlan: def _module( task_name: str | None = "traj_arm", + physical_operation_timeout: float = 60.0, + gateway: FakeCoordinatorGateway | None = None, ) -> tuple[ManipulationModule, FakeCoordinatorGateway]: module = make_module() + module.config.physical_operation_timeout = physical_operation_timeout config = _config(task_name) module._robots = {"arm": ("robot_id", config, object())} module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([config]) - gateway = install_runtime(module, [config]) - return module, gateway + return module, install_runtime(module, [config], gateway) @pytest.fixture(autouse=True) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index a71536f2b9..00313c632b 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -958,6 +958,8 @@ def _apply_target_evaluation_result( or request.group_ids != self.state.selected_group_ids ): return + if self.state.manipulation_state == "FAULT": + return collision_free = result.collision_free success = result.success self.state.feasibility.status = self._feasibility_status(result, success, collision_free) @@ -1319,6 +1321,8 @@ def _finish_operation( and not self._operation_is_current(operation_id, selection_epoch) ): return + if self.state.manipulation_state == "FAULT": + return self.state.action_status = ActionStatus.IDLE if clear_error: self.state.error = "" @@ -1335,6 +1339,8 @@ def _set_recoverable_error(self, message: str) -> None: if self._closed: return self.refresh() + if self.state.manipulation_state == "FAULT": + return # This is a panel-local validation message, not a runtime diagnostic. self.state.error = message self._update_status_text() diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index 4e8090f229..dc2e506536 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -171,10 +171,6 @@ def can_cancel(self) -> bool: "DISPATCHING", "RUNNING", "CANCELLING", - } or self.action_status in { - ActionStatus.RUNNING, - ActionStatus.PREVIEWING, - ActionStatus.EXECUTING, } def can_execute(self, action_status: ActionStatus | None = None) -> bool: diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index 0116821a81..c0a85d5465 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -33,6 +33,7 @@ OperationWorker, PanelRuntime, PlanStatus, + TargetEvaluationRequest, TargetEvaluationWorker, TargetStatus, ) @@ -524,3 +525,38 @@ def test_gui_projects_all_runtime_states_and_fault_reset_preserves_local_failure assert gui.state.action_status == ActionStatus.FAILED assert gui.state.local_action_failure is True + + +def test_late_local_callbacks_cannot_clear_runtime_fault() -> None: + gui = make_gui() + gui.state.manipulation_state = "FAULT" + gui.state.action_status = ActionStatus.FAILED + gui.state.runtime_failure = True + gui.state.error = "runtime fault" + gui.state.selected_group_ids = ("arm/manipulator",) + + gui._finish_operation("preview=True") + gui._apply_target_evaluation_result( + TargetEvaluationRequest(0, "joints", group_ids=("arm/manipulator",)), + TargetEvaluationResult(True, "FEASIBLE", "ok", True), + ) + + assert gui.state.manipulation_state == "FAULT" + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "runtime fault" + assert gui.state.can_cancel() is False + + +def test_stale_execute_validation_cannot_replace_runtime_fault_diagnostic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + gui = make_gui() + gui._apply_operator_status(OperatorStatus("FAULT", diagnostic="hardware fault")) + monkeypatch.setattr(gui, "refresh", lambda: None) + + gui._submit_execute() + + assert gui.state.manipulation_state == "FAULT" + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "hardware fault" + assert gui.state.can_execute() is False diff --git a/dimos/manipulation/visualization/viser/test_state.py b/dimos/manipulation/visualization/viser/test_state.py index a9ec709b93..ee49c8332b 100644 --- a/dimos/manipulation/visualization/viser/test_state.py +++ b/dimos/manipulation/visualization/viser/test_state.py @@ -74,12 +74,20 @@ def test_selection_epoch_change_resets_plan_and_invalidates_sequence() -> None: assert state.plan_state.status == PlanStatus.NONE -def test_cancel_is_available_for_running_preview_and_execute_actions() -> None: +def test_cancel_uses_only_the_runtime_lifecycle_gate() -> None: state = PanelState() - for action in (ActionStatus.RUNNING, ActionStatus.PREVIEWING, ActionStatus.EXECUTING): - state.action_status = action - assert state.can_cancel() is True + for lifecycle in ("PLANNING", "DISPATCHING", "RUNNING", "CANCELLING"): + state.manipulation_state = lifecycle + for action in ActionStatus: + state.action_status = action + assert state.can_cancel() is True + + for lifecycle in ("IDLE", "READY", "FAULT"): + state.manipulation_state = lifecycle + for action in (ActionStatus.RUNNING, ActionStatus.PREVIEWING, ActionStatus.EXECUTING): + state.action_status = action + assert state.can_cancel() is False def test_operation_worker_uses_per_operation_timeout() -> None: diff --git a/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml b/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml new file mode 100644 index 0000000000..c0a8162549 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-21 diff --git a/openspec/changes/simplify-manipulation-execution-runtime/design.md b/openspec/changes/simplify-manipulation-execution-runtime/design.md new file mode 100644 index 0000000000..2e28874d8d --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/design.md @@ -0,0 +1,111 @@ +# Design: Simplified Manipulation Execution Runtime + +## Context + +The just-added manipulation runtime is approximately 2,084 LOC. It currently spreads coordination across a large reducer, poller and action-specific threads, deadline/reset/gateway-close threads, duplicated flags and histories, and module-side timeout reconciliation. That structure obscures the single-owner lifecycle contract even though the required behavior is already defined and must remain unchanged. + +The refactor is an internal, behavior-preserving rewrite with a compact initial-runtime scope. The public `ManipulationModule` API and all safety outcomes remain the compatibility boundary. Blocking coordinator RPCs must not block the owner loop, while lifecycle state must still have one writer and deterministic ordering. + +## Goals / Non-Goals + +### Goals + +- Preserve four explicit seams and the achieved policy/effects/auxiliary/passive module separations: + - plan materialization: 220–280 LOC; + - coordinator gateway: 80–120 LOC; + - models: 100–150 LOC in the final simplified architecture; + - runtime owner: no hard LOC threshold; keep responsibilities reviewable without splitting a required public façade from its private owner. +- Preserve all existing public `ManipulationModule` semantics, including serialized lifecycle and ready-plan behavior. +- Preserve minimal operation/task/method RPC correlation, sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical completion polling, UNKNOWN fail-closed safety faults, and truthful shutdown. +- Keep lifecycle state and mutable ownership in one runtime owner loop. +- Make blocking RPC completion explicit through futures and result waits. +- Enforce a module-configured `physical_operation_timeout: float = 60.0` safety deadline for operations that remain `RUNNING` indefinitely. +- Accept an additive runtime `monotonic_clock` seam defaulting to `time.monotonic` for deterministic deadline behavior. +- Keep the initial runtime compact: retryable reset generations, rich action/reset generations, exact result-history depth, and eventual post-deadline shutdown success are deferred rather than implemented promises. +- Define success by reduced duplicated responsibility, tested policy/effect/auxiliary/passive boundaries, and no observable behavior regression; LOC is an observation, not a release gate. + +### Non-Goals + +- No redesign of planning algorithms, trajectory data, coordinator protocol, or public manipulation semantics. +- No freshness-based or robot-name-based execution selection; neither is a replacement safety gate. +- No new poller, per-action deadline, reset, gateway-close, or module-side timeout-reconciliation thread. +- No UI/layout work, unrelated source cleanup, or changes outside the requested change during implementation. + +## Decisions + +### Four narrow implementation seams + +Plan materialization converts planner output into a complete executable plan and owns no lifecycle state. Models contain immutable records/enums/results and the single mutable owner-state record. The gateway is a thin coordinator RPC adapter. The actor owns lifecycle handlers, minimal correlation, dispatch sequencing, polling, conditional compensation, and shutdown. + +The alternative of retaining the existing reducer and extracting helpers was rejected: it would preserve duplicate state representations and make ownership and transition behavior difficult to verify. The alternative of splitting into many event-specific services was rejected because it would recreate cross-thread ordering problems. + +### Fixed daemon blocking-RPC executor + +All potentially blocking coordinator RPCs run on exactly four daemon executor workers. Each submission registers operation, task, and method correlation before the effect and returns a future; the owner loop consumes completed results at defined handler boundaries. RPC deadlines are represented as future/result outcomes and safety faults, not as per-action timer threads. Unresolved work blocks replacement execution; no rich action/reset generation model is required. + +The executor is fixed for the runtime lifetime and is shut down truthfully: one shutdown initiator seals admission, shutdown stops accepting work, observer-only waits observe known futures, and incomplete/unknown remote work is reported rather than pretending it is complete. Concurrent shutdown callers are rejected. A shutdown deadline does not promise eventual success after the deadline; daemon workers prevent process teardown from being held hostage, but do not weaken fail-closed outcomes. + +Admission has one exception to strict owner-only writes: a caller-owned latch may seal gateway admission before the owner processes shutdown. That latch is an admission gate only; all `RuntimeContext` writes, including shutdown transitions, remain owner-loop writes. + +### Owner-loop timer scan + +One daemon owner loop serializes requests, completed future results, and a monotonic timer scan. The owner timer is authoritative for physical-completion polling, deadline observation, cancellation reconciliation, deferred reset handling, and shutdown observation. It does not perform RPCs itself; it schedules blocking work on the fixed executor and applies results only after correlation checks. + +This replaces poller, per-action deadline, reset, and gateway-close threads. A single timer scan is preferred over independent timers because it preserves ordering and avoids races between timeout, cancellation, reset, and shutdown. + +### Monotonic clock seam and public polling + +The runtime constructor keeps a compatible callable injection named `monotonic_clock`, defaulting to `time.monotonic`. It wraps that callable in a dedicated `ValidatedMonotonicClock` source wrapper. The wrapper serializes access, validates that every sample is finite and nondecreasing, and is the only clock exposed to runtime deadline code. An invalid sample violates the injected dependency contract and raises a dependency-contract error; it is not converted into lifecycle `FAULT` state or a remote safety outcome. + +Every deadline creation, deadline comparison, and relative wait calculation routes through `ValidatedMonotonicClock`. Wall-clock time, mixed clock domains, and ad hoc elapsed-time calculations are not allowed for runtime deadlines. The compatible constructor surface remains callable injection/default behavior; validation and serialization are supplied by the wrapper rather than by changing callers. + +The owner timer is the authoritative deadline processor. Public immediate polling is supported only while an operation is `RUNNING`; it does not provide a special narrow timer-precheck guarantee and does not promise `STATUS` in reset, shutdown, gateway, or other non-running states. + +### Focused handlers and one mutable state + +Phase 4 is decomposed into reviewable slices. Slice 1 adds the `_set_state` write +wrapper and extracts only planning, dispatch-admission, and lifecycle-deadline +branches. It intentionally retains the immutable context representation and all +Phase 3 scheduler/executor mechanics. Later slices will migrate state ownership, +remove redundant bookkeeping, and extract task/cancellation/reset/shutdown +handlers. + +Handlers are grouped by concern: planning/ready-plan materialization, execute admission, sequential task dispatch, task-result reconciliation, physical-operation deadline expiry, conditional cancellation/compensation, deferred reset handling, and shutdown. Each handler mutates only the owner-owned state record and publishes the existing snapshot/result contract. Correlation IDs are limited to operation/task/method RPC identity; stale or unresolved results fail closed and cannot replace active work. Result retention is small and current/recent, with no exact history-depth promise. + +The giant reducer, redundant flags, history/alias fields, and module-side timeout reconciliation are removed. Public methods remain adapters that submit owner-loop requests and wait for the corresponding future/result, rather than independently reconciling lifecycle state. This does not require a separate public façade/private owner split: the required boundary is behavioral ownership and tested effects, not a numeric actor size or an additional public layer. + +The physical-operation deadline starts when the operation enters `RUNNING` and is 60 seconds by default. It is not planning time, an action/RPC deadline, a caller observation timeout, or trajectory duration. On expiry the owner schedules cancellation and reconciliation; only authoritative reconciliation can yield completion, otherwise the owner publishes safe terminal `FAILED` or a correlated `FAULT`. The module never reconciles this deadline. + +### Sequential multi-task dispatch and compensation + +The actor dispatches the complete plan one coordinator task at a time. It registers operation, task, and method correlation before each RPC. If a later task fails after an earlier task is known or potentially active, compensation is invoked in order; if activity is UNKNOWN, the outcome remains fail-closed and replacement is blocked. No second operation is admitted while unresolved work remains. + +### Reset and shutdown scope + +Reset generations that can safely be retried are deferred to a later scope. An unsafe or timed-out reset remains `FAULT` and is not retried in-process. Shutdown has one initiator; concurrent callers are rejected, observer-only waits do not mutate lifecycle state, and expiry reports unresolved work without promising eventual success. Admission sealing is the sole caller-side latch exception before owner shutdown processing. + +## Risks / Trade-offs + +- **Blocking RPC exceeds its expected duration** → represent the result as an explicit deadline/unknown safety outcome, let the owner loop continue its timer scan, and preserve reset proof requirements. +- **A stale or unresolved future completes after reset or shutdown** → retain minimal operation/task/method correlation until classification; never apply it to new work, and fail closed when activity is unknown. +- **Daemon executor hides unfinished work at process exit** → truthful shutdown reports unresolved work and does not claim cancellation or completion without proof. +- **LOC targets encourage over-compression** → treat measurements as maintainability evidence only; prioritize reduced duplicated responsibility, tested boundaries, and the required safety matrix over an arbitrary actor-size gate. +- **Phase 1 faithful extraction exceeds the final models budget** → retain the maintainable 190-LOC models baseline during pure extraction; defer reduction to the later simplification phases and do not compress records unsafely. +- **Changing test seams masks behavior changes** → run characterization and black-box compatibility tests before and after each extraction; decouple white-box tests only after equivalent observable coverage exists. +- **A clock seam hides an unsafe clock** → put finite/nondecreasing validation and serialized access in `ValidatedMonotonicClock`; treat invalid samples as dependency-contract failures, not lifecycle faults, and test the wrapper independently. +- **Public polling races the timer owner** → make the owner timer authoritative and limit immediate public polling to `RUNNING` operations. + +## Migration Plan + +1. Freeze characterization coverage and extract pure modules with re-exports; no behavior change. +2. Introduce futures/result waits behind the existing gateway and verify correlation and outcomes. +3. Replace thread topology with the fixed executor and owner-loop timer scan. +4. Move logic into focused handlers and one mutable owner state, deleting redundant reducer data. +5. Clean module/API adapters while preserving public semantics and removing freshness/robot-name selection paths. +6. Keep only compact seam tests and the required safety matrix; defer rich reset-generation/history coverage. + +Rollback at each phase is a file-level revert before the next phase; no data migration or deployment migration is required. + +## Open Questions + +- None blocking implementation. Existing polling cadence, RPC deadline values, and reset proof criteria remain authoritative and must be carried over unchanged. diff --git a/openspec/changes/simplify-manipulation-execution-runtime/proposal.md b/openspec/changes/simplify-manipulation-execution-runtime/proposal.md new file mode 100644 index 0000000000..3e153556c3 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/proposal.md @@ -0,0 +1,35 @@ +## Why + +The newly added manipulation execution runtime preserves the required safety and lifecycle behavior, but its 2,084 lines concentrate planning materialization, RPC adaptation, mutable lifecycle state, polling, deadlines, reset handling, and shutdown mechanics in one over-engineered implementation. Simplifying those seams now will make the behavior easier to audit and maintain without changing the public `ManipulationModule` contract or safety semantics. + +## What Changes + +- Retain and document the achieved separations between policy, effects, auxiliary, and passive modules, with explicit boundaries for plan materialization, the coordinator gateway, runtime models, and the runtime owner. +- Replace per-action/poller/deadline/reset/gateway-close threads with one fixed daemon blocking-RPC executor and one owner-loop timer scan. +- Replace the giant reducer and redundant flags/history/aliases with focused handlers over one owner-owned mutable runtime state. +- Use futures and observer-only waits for blocking coordinator operations, preserving minimal operation/task/method correlation and serialized lifecycle decisions. Unresolved work blocks replacement execution. +- Enforce a runtime-owned `physical_operation_timeout: float = 60.0` safety deadline for operations that remain `RUNNING`, with owner-scheduled cancellation/reconciliation and safe terminal failure or correlated fault. +- Enforce strict owner-only `RuntimeContext` writes; permit only a caller-owned admission latch to seal gateway admission before owner shutdown processing. +- Use a dedicated serialized `ValidatedMonotonicClock` wrapper behind compatible `monotonic_clock` callable injection/default behavior; reject invalid samples as dependency-contract failures rather than lifecycle faults. +- Preserve sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical-completion polling, UNKNOWN fail-closed safety outcomes, the 60-second physical deadline, admission sealing, and truthful shutdown. +- Defer retryable reset generations; unsafe or timed-out reset remains `FAULT` with no in-process retry. Retain only a small current/recent result set rather than promising exact history depth. +- Permit one shutdown initiator only: concurrent shutdown callers are rejected, and no eventual success is promised after the shutdown deadline. +- Make the owner timer authoritative. Public immediate polling is only for `RUNNING` operations and has no special narrow timer-precheck guarantee. +- **BREAKING** Remove implementation-only freshness/robot-name execution selection paths; retain the existing public manipulation semantics without freshness- or robot-name-based execution selection. +- Keep all changes confined to the new internal runtime structure; do not add execution selection based on freshness or robot name. +- Do not impose a hard runtime-actor LOC threshold or require a public-façade/private-owner split; judge the refactor by reduced duplicated responsibility, tested boundaries, and unchanged behavior. + +## Capabilities + +### New Capabilities +- `manipulation-lifecycle-runtime`: Compact fail-closed lifecycle ownership, minimal correlation, safety, reset, and shutdown semantics for the simplified runtime. +- `coordinator-trajectory-execution`: Compact gateway, sequential dispatch, conditional compensation, polling, and blocking-RPC execution semantics. + +### Modified Capabilities + +## Impact + +- Affects only the implementation and tests of the recently added manipulation execution runtime and its `ManipulationModule` integration. +- Introduces no external dependency and no source changes outside the eventual runtime refactor. +- Preserves public `ManipulationModule` methods, lifecycle outcomes, serialized ready-plan behavior, and operator-visible safety results. +- Success is reduced duplicated responsibility across the existing policy/effects/auxiliary/passive separations, tested ownership and effect boundaries, and no behavior regression in the safety and lifecycle contract. diff --git a/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md b/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md new file mode 100644 index 0000000000..6aca794f0a --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md @@ -0,0 +1,51 @@ +## ADDED Requirements + +### Requirement: Use a fixed blocking-RPC executor and owner timer scan +The runtime SHALL execute blocking coordinator RPCs through one fixed executor with exactly four daemon workers returning futures. One owner-loop timer scan SHALL be authoritative for polling, deadline observation, cancellation reconciliation, deferred reset handling, and shutdown observation. The runtime SHALL NOT create poller, per-action deadline, reset, gateway-close, or module-side timeout-reconciliation threads. + +The runtime SHALL include the module-configured `physical_operation_timeout: float = 60.0` safety deadline for a coordinator state that remains `RUNNING` indefinitely. Its expiry SHALL be handled by the owner loop, which schedules cancellation/reconciliation and reports safe terminal `FAILED` or a correlated `FAULT`; it SHALL not be confused with planning, per-RPC action, caller wait, or trajectory-duration timeouts. + +The runtime constructor SHALL retain compatible callable injection through `monotonic_clock`, defaulting to `time.monotonic`, and SHALL route it through a dedicated serialized `ValidatedMonotonicClock` wrapper. The wrapper SHALL validate finite, nondecreasing samples and SHALL be the sole source used for deadline creation, deadline comparison, and relative wait calculations. An invalid sample SHALL fail the clock dependency contract, not create lifecycle `FAULT` state. + +#### Scenario: RPC does not block lifecycle ownership +- **WHEN** a coordinator RPC blocks +- **THEN** the RPC SHALL run on the fixed executor +- **AND** the owner loop SHALL continue processing its serialized timer/request path + +#### Scenario: Indefinite running state reaches the physical deadline +- **WHEN** a coordinator task remains `RUNNING` beyond `physical_operation_timeout` +- **THEN** the owner loop SHALL enqueue cancellation and reconciliation +- **AND** the eventual result SHALL be terminal `FAILED` or a correlated `FAULT` +- **AND** duration alone SHALL never be treated as physical completion + +### Requirement: Preserve sequential dispatch and compensation +The gateway/actor boundary SHALL preserve sequential multi-task dispatch for complete plans. It SHALL register minimal operation, task, and method correlation before each RPC effect. It SHALL invoke compensation only when earlier tasks are known or potentially active; UNKNOWN activity SHALL fail closed and block replacement execution. + +#### Scenario: Later task failure compensates earlier work +- **WHEN** a multi-task plan establishes an earlier task and a later task is rejected or becomes uncertain +- **THEN** the runtime SHALL preserve the existing compensation sequence +- **AND** the normalized operation result SHALL reflect compensation success or fail-closed safety uncertainty + +### Requirement: Preserve physical completion polling +Dispatch acceptance SHALL remain distinct from physical completion. The timer scan SHALL poll active coordinator tasks and transition to the existing terminal outcome only after authoritative completion or safe cancellation/inactivity proof. + +#### Scenario: Accepted task remains active +- **WHEN** a coordinator accepts a task without reporting physical completion +- **THEN** the runtime SHALL retain the active correlated task +- **AND** it SHALL report the existing running state rather than completion + +#### Scenario: Immediate public polling is running-only +- **WHEN** public `poll` is called +- **THEN** it SHALL issue immediate physical polling only for a `RUNNING` operation +- **AND** the owner timer SHALL process deadlines and other timer work +- **AND** `STATUS` SHALL not be guaranteed in reset, shutdown, gateway, or other non-running states + +### Requirement: Keep the gateway thin and synchronous at its boundary +The coordinator gateway SHALL contain only request/result adaptation and RPC invocation required by the actor. It SHALL not own lifecycle flags, timers, retries, or independent close state; futures and observer-only waits SHALL carry completion back to the owner loop. Unresolved work SHALL block replacement execution. + +The policy, effects, auxiliary, and passive module boundaries SHALL remain explicit and testable. The runtime SHALL NOT require a numeric actor-size limit or an additional public-façade/private-owner layer; behavioral ownership and effect isolation are the contract. + +#### Scenario: Gateway result is applied by owner +- **WHEN** the gateway future completes with acceptance, rejection, cancellation, completion, or uncertainty +- **THEN** the owner loop SHALL normalize and apply the result using minimal operation/task/method correlation +- **AND** the gateway SHALL not publish an independent lifecycle state diff --git a/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md b/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md new file mode 100644 index 0000000000..9447383be6 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md @@ -0,0 +1,71 @@ +## ADDED Requirements + +### Requirement: Preserve serialized lifecycle ownership +The simplified runtime SHALL preserve the existing public `ManipulationModule` lifecycle semantics and SHALL have exactly one owner loop that mutates lifecycle, ready-plan, operation, task, fault, and shutdown state. Any attempt/action/reset bookkeeping is compact local correlation, not a rich generation model. + +The owner SHALL be the only writer of `RuntimeContext` and its mutable state. A caller-owned admission latch MAY be written outside the owner only to seal gateway admission before the owner processes shutdown; it SHALL not mutate lifecycle, operation, task, fault, or shutdown state. + +The runtime SHALL use one owner loop, exactly four daemon RPC workers, and observer-only caller waits. It SHALL retain only a small current/recent result set; no exact result-history depth is part of the contract. + +The policy, effects, auxiliary, and passive module separations SHALL remain explicit and testable. Success SHALL be assessed by reduced duplicated responsibility, tested ownership/effect boundaries, and no observable behavior regression, not by a hard actor LOC threshold or a required public-façade/private-owner split. + +The runtime constructor SHALL retain compatible callable injection through `monotonic_clock`, defaulting to `time.monotonic`, and SHALL wrap it in a dedicated `ValidatedMonotonicClock`. The wrapper SHALL serialize source access and SHALL accept only finite, nondecreasing samples. All runtime deadline creation, deadline comparison, and relative wait calculations SHALL route through the wrapper. An invalid sample SHALL be a dependency-contract failure, not a lifecycle `FAULT`, and SHALL not be converted into a remote safety outcome. + +#### Scenario: Concurrent lifecycle requests are serialized +- **WHEN** execute, cancel, reset, planning results, and timer events arrive concurrently +- **THEN** the owner loop SHALL apply them in one deterministic order +- **AND** no caller-visible snapshot SHALL combine state from different operations + +### Requirement: Preserve correlation and normalized outcomes +The runtime SHALL retain minimal operation, task, and method correlation through planning, sequential dispatch, polling, cancellation, compensation, deadlines, faults, and shutdown. Registration SHALL occur before each RPC effect. Stale, mismatched, or unresolved results SHALL NOT mutate the current operation, and unresolved work SHALL block replacement execution; rich action/reset generations are out of scope. + +#### Scenario: Late result cannot affect a new operation +- **WHEN** an RPC future from a prior operation completes after reset or a new operation begins +- **THEN** the owner loop SHALL classify it using its original correlation +- **AND** it SHALL not apply the result to the new operation + +### Requirement: Preserve safety and shutdown semantics +Unknown, malformed, unavailable, contradictory, or deadline-uncertain remote results SHALL fail closed and preserve the existing safety-fault behavior. Unsafe or timed-out reset SHALL remain `FAULT` with no in-process retry; retryable reset generations SHALL be deferred. Shutdown SHALL report unresolved work truthfully. + +#### Scenario: Deadline uncertainty faults safely +- **WHEN** a blocking coordinator result cannot be obtained before its established deadline +- **THEN** the runtime SHALL preserve the existing safety fault outcome +- **AND** it SHALL not report completion or safe reset without proof + +#### Scenario: Unsafe reset remains faulted +- **WHEN** reset is requested while remote inactivity is active, uncertain, or timed out +- **THEN** the runtime SHALL retain the fault/safety condition and issue no new execution +- **AND** it SHALL not retry the reset generation in-process + +#### Scenario: Shutdown is truthful +- **WHEN** shutdown begins with an unresolved RPC or remote attempt +- **THEN** the runtime SHALL not report that work as completed or cancelled without proof +- **AND** the public shutdown result SHALL reflect the unresolved condition + +#### Scenario: Shutdown has one initiator +- **WHEN** a shutdown is already being processed by one caller +- **THEN** a concurrent shutdown caller SHALL be rejected +- **AND** expiry of the shutdown deadline SHALL report unresolved work without promising eventual success + +### Requirement: Enforce a runtime-owned physical operation deadline +Each module-configured runtime SHALL expose `physical_operation_timeout: float = 60.0`. The owner SHALL start this safety deadline when a physical operation enters `RUNNING` and SHALL owner-schedule cancellation and reconciliation when it expires. This deadline is distinct from `planning_timeout`, per-RPC `action_timeout`, caller observation wait timeouts, and trajectory duration. + +#### Scenario: Physical operation exceeds its safety deadline +- **WHEN** a correlated operation remains `RUNNING` until `physical_operation_timeout` expires +- **THEN** the owner SHALL schedule cancellation and reconciliation +- **AND** it SHALL produce a safe terminal `FAILED` result or a correlated `FAULT` +- **AND** it SHALL not infer completion from trajectory duration or permit module-side reconciliation + +#### Scenario: Immediate polling is running-only +- **WHEN** the public `poll` method is called +- **THEN** it SHALL perform immediate physical polling only for a `RUNNING` operation +- **AND** the owner timer SHALL remain authoritative for deadlines, reset, shutdown, and other timer processing +- **AND** no special precheck SHALL guarantee `STATUS` during reset, shutdown, gateway, or other non-running states + +### Requirement: No execution selection by freshness or robot name +The simplified runtime SHALL preserve the contract that execution is not selected by freshness or robot name. It SHALL dispatch the complete plan represented by the planning groups and SHALL reject partial or robot-name-selected execution without issuing coordinator work. + +#### Scenario: Complete plan dispatch ignores freshness selection +- **WHEN** a complete ready plan reaches execution with a non-fresh manipulation-layer observation +- **THEN** the runtime SHALL not reject it solely for that observation +- **AND** it SHALL continue through normal coordinator acceptance and safety handling diff --git a/openspec/changes/simplify-manipulation-execution-runtime/tasks.md b/openspec/changes/simplify-manipulation-execution-runtime/tasks.md new file mode 100644 index 0000000000..7469cabc28 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/tasks.md @@ -0,0 +1,71 @@ +## 1. Pure extraction and re-exports + +Baseline before Phase 1: `execution_runtime.py` was 2,084 LOC; topology/materialization occupied lines 82–337 and public models/results occupied lines 342–469. The public boundary was the `dimos.manipulation.execution_runtime` import path used by the manipulation module, tests, and helpers. Preserve the achieved policy/effects/auxiliary/passive module separations; LOC is recorded only as maintainability evidence, not as a hard actor gate or a requirement for a public-façade/private-owner split. + +- [x] 1.1 Record current public `ManipulationModule` behavior and LOC/boundary baselines without changing source behavior. +- [x] 1.2 Extract pure plan materialization into a 220–280 LOC internal module with no lifecycle ownership. +- [x] 1.3 Extract a faithful models and correlation/result baseline into an internal module (190 LOC in Phase 1; the 100–150 LOC target remains a later simplification target, not a reason for unsafe compression). +- [x] 1.4 Add compatibility re-exports so existing runtime imports and public module semantics remain unchanged. +- [x] 1.5 **Validation gate:** run plan/materialization, model, import, and existing manipulation characterization tests; no behavior changes are allowed. + +## 2. Futures and explicit result waits + +- [x] 2.1 Extract the coordinator gateway into an 80–120 LOC adapter containing only RPC request/result conversion. +- [x] 2.2 Route blocking coordinator calls through retained futures while registering minimal operation/task/method correlation before each RPC effect. The existing auxiliary-thread topology remains intentionally unchanged until Phase 3. +- [x] 2.3 Replace action-thread completion signaling at the owner boundary with explicit correlated completion results and observer-only waits with normalized outcomes. +- [x] 2.4 **Validation gate:** verify immediate acceptance, rejection, completion, cancellation, malformed, unavailable, and stale-result cases against characterization tests. + +## 3. Fixed executor and owner-loop timer scan + +- [x] 3.1 Add one fixed blocking-RPC executor with exactly four daemon workers, bounded lifetime, and truthful unresolved-work shutdown. +- [x] 3.2 Add one owner-loop timer scan for polling, deadlines, cancellation reconciliation, deferred reset handling, and shutdown observation. +- [x] 3.3 Remove poller, per-action deadline, reset, and gateway-close threads without changing polling cadence or deadline semantics. +- [x] 3.4 Keep module-side timeout reconciliation out of the new path. +- [x] 3.5 **Validation gate:** run concurrency, deadline/UNKNOWN fail-closed, physical-completion, deferred-reset, and truthful-shutdown tests; inspect four-worker ownership. +- [x] 3.6 Add the module-configured `physical_operation_timeout: float = 60.0` safety deadline for indefinitely `RUNNING` coordinator state; owner-schedule cancellation/reconciliation on expiry and produce safe terminal `FAILED` or correlated `FAULT`. Keep it distinct from planning, per-RPC action, caller observation, and trajectory-duration timeouts. +- [x] 3.7 Add the additive compatible `monotonic_clock` callable injection, defaulting to `time.monotonic`, and route it through a dedicated serialized `ValidatedMonotonicClock` wrapper; validate finite/nondecreasing samples, treat invalid samples as dependency-contract failures rather than lifecycle `FAULT`, and route every deadline creation/comparison and relative wait calculation through the wrapper. + +## 4. Focused handlers and mutable owner state + +- [x] 4.1a Add a behavior-preserving `_set_state(**changes)` owner-state wrapper. +- [x] 4.1b Extract planning lifecycle transitions into focused owner-only handlers. +- [x] 4.1c Enforce owner-only `RuntimeContext` writes in dispatch admission and lifecycle transitions; permit only a caller-owned latch that seals gateway admission before owner shutdown processing. +- [ ] 4.1d **Deferred/out of scope:** rich handler decomposition for reset generations and action-level transitions; retain only the compact task dispatch, result reconciliation, deadline, conditional compensation, and shutdown seam. +- [x] 4.2a Introduce the mutable owner-state record. +- [x] 4.2b Migrate lifecycle state writes while preserving immutable snapshots and strict owner-only `RuntimeContext` ownership. +- [x] 4.2c Migrate minimal operation/task/method correlation state writes under the owner-only rule. +- [ ] 4.3 **Deferred/replaced:** exact history/alias removal and rich action/reset bookkeeping are out of scope; retain only small current/recent results and the existing compact policy/effects/auxiliary/passive boundaries. +- [x] 4.4 Preserve sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical completion polling, and safe rejection of overlapping execution. +- [x] 4.5a Verify planning and ready-plan replacement behavior. +- [ ] 4.5b **Deferred/replaced:** rich dispatch-admission and single-use-plan bookkeeping is out of scope; verify only unresolved-work admission blocking. +- [x] 4.5c Verify the required compact safety matrix: UNKNOWN fail-closed, sequential dispatch, conditional compensation, 60-second physical deadline, cancellation, unsafe/timed-out reset remaining `FAULT`, and truthful shutdown. +- [x] 4.5e Verify immediate public polling is `RUNNING`-only; no special narrow timer-precheck or `STATUS` guarantee applies in reset/shutdown/gateway states. +- [x] 4.5d **Validation gate:** run the required compact Phase-4 safety matrix after all slices. + +## 5. Module and API cleanup + +- [x] 5.1 Keep `ManipulationModule` as a public adapter over the runtime and remove module-owned lifecycle mutation and reconciliation. +- [x] 5.2 Remove freshness- and robot-name-based execution selection and partial-plan paths without adding a replacement selection gate. +- [x] 5.3 Preserve all existing public `ManipulationModule` semantics, status snapshots, normalized outcomes, and truthful shutdown behavior. +- [x] 5.4 **Validation gate:** run module, operator, UI/state-adapter, API, and compatibility tests; confirm no unrelated source paths changed. + +## 6. White-box test decoupling + +- [x] 6.1 Move white-box tests from reducer/thread internals to plan, model, gateway, actor-handler, and owner-loop seams. +- [x] 6.2 Add focused tests for minimal correlation isolation, future completion ordering, owner timer scans, conditional compensation, deferred reset handling, executor shutdown, and the monotonic clock seam. +- [x] 6.2a Add deterministic tests for stale and cleared physical deadlines, including nondecreasing clock samples and no cancellation when the deadline is absent or not due. +- [x] 6.2b Add deterministic tests that retained faults remain observable after physical-deadline handling and that expiry yields safe `FAILED` or correlated `FAULT`. +- [x] 6.2c Add public poll tests proving immediate polling is `RUNNING`-only and that `STATUS` is not guaranteed in reset/shutdown/gateway states. +- [x] 6.2d Add poll-versus-shutdown tests covering clock-driven physical expiry, admission sealing, one shutdown initiator, concurrent rejection, and truthful unresolved shutdown without eventual-success promise. +- [x] 6.2e Add isolated `ValidatedMonotonicClock` unit tests for default/injected callables, serialized access, finite-sample rejection, nondecreasing-sample rejection, and dependency-contract error classification without lifecycle mutation. +- [x] 6.2f Add deterministic lifecycle tests proving all deadline and relative-wait paths use the wrapper, including stale/cleared deadlines, retained faults, physical expiry, deferred reset behavior, and state-dependent poll/shutdown behavior. +- [x] 6.3 Retain black-box characterization coverage for every public lifecycle and safety contract. +- [x] 6.4 **Validation gate:** run targeted manipulation runtime and coordinator test suites with no skipped contract scenarios. + +## 7. Final verification + +- [x] 7.1 Record the compact seam measurement as maintainability evidence and verify reduced duplicated responsibility across the policy/effects/auxiliary/passive boundaries; there is no hard actor LOC gate or required public-façade/private-owner split. +- [x] 7.2 Verify strict owner-only `RuntimeContext` writes, the sole admission-latch exception, exactly four daemon RPC workers, absence of forbidden thread types, and absence of module-side timeout reconciliation. +- [x] 7.3 Verify physical-operation deadline expiry schedules owner cancellation/reconciliation and yields safe terminal `FAILED` or correlated `FAULT`, never duration-inferred completion. +- [x] 7.4 Run formatting, lint, type checks, targeted manipulation/coordinator tests, the full repository test, and OpenSpec strict validation; formatting/lint/type checks, the manipulation target suite (364 passed, 10 skipped, 16 deselected), OpenSpec validation, and diff review passed. The full repository suite was environment-limited by an unrelated MCP `localhost:9990` port conflict/connection resets (2910 passed, 30 skipped, 289 deselected, 2 failed, 24 errors). +- [x] 7.5 Review the final diff and task ledger so no source implementation task is marked complete merely because work was planned, and only supported completed work remains checked.