From 9629eb74e3208f30be60c8486dce6bfa4744930b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 10:43:20 +0000 Subject: [PATCH 1/7] docs: add cuRobo V2 integration design --- .gitignore | 4 + .../plans/2026-07-11-curobo-v2-integration.md | 846 ++++++++++++++++++ ...2026-07-11-curobo-v2-integration-design.md | 282 ++++++ 3 files changed, 1132 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-11-curobo-v2-integration.md create mode 100644 docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md diff --git a/.gitignore b/.gitignore index f84991ad..0064cc2d 100644 --- a/.gitignore +++ b/.gitignore @@ -203,3 +203,7 @@ embodichain/VERSION # benchmark results scripts/benchmark/rl/reports/* + +# Local agent execution state and isolated worktrees +.superpowers/ +.worktrees/ diff --git a/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md b/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md new file mode 100644 index 00000000..27480f21 --- /dev/null +++ b/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md @@ -0,0 +1,846 @@ +# CuRobo V2 Motion-Planning Integration Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add NVIDIA cuRobo V2 as an optional collision-aware EmbodiChain planner, route it through `MotionGenerator` and supported atomic actions, verify it at unit/CUDA/DexSim levels, and ship a runnable Panda obstacle-avoidance demo. + +**Architecture:** `CuroboPlanner(BasePlanner)` owns lazy V2 bindings and converts standard batched `PlanState` inputs into V2 `JointState` and `GoalToolPose` calls. `MotionGenerator` becomes capability-driven rather than special-casing `NeuralPlanner`; atomic actions keep their existing full-DoF output contract and select the CuRobo path only when configured. The backend returns collision-checked cuRobo samples without EmbodiChain IK pre-processing or unsafe post-resampling. + +**Tech Stack:** Python 3.10+, PyTorch, EmbodiChain `@configclass`, DexSim, pytest, NVIDIA cuRobo V2 (tested against upstream `3fd54dc782a82e5500a771cfd47856ea499d5fef`), CUDA. + +## Global Constraints + +- Support cuRobo V2 only; never import, install, or fallback to V1 `curobo.wrap.reacher.motion_gen`. +- cuRobo is optional: `import embodichain.lab.sim.planners` must work without the `curobo` module; import V2 only inside a lazy helper used by `CuroboPlanner`. +- Core dependencies in `pyproject.toml` remain unchanged. Installation instructions use NVIDIA's matching `.[cu12]`, `.[cu13]`, `.[cu12-torch]`, or `.[cu13-torch]` extras. +- Use Apache 2.0 headers, `from __future__ import annotations`, full public type annotations, Google-style docstrings, `@configclass` for configuration classes, and `__all__` in every new public Python module. +- Preserve `PlanState` and `PlanResult` tensors: leading batch dimension `B`; positions/velocities/accelerations `(B, N, D)`; `dt` `(B, N)`; duration `(B,)`; success `(B,)`. +- The first release supports one configured control part per request. It rejects coordinated dual-arm CuRobo requests and does not add attachment/detachment collision modeling or ActionBank support. +- The scene boundary accepts only cuRobo-supported `cuboid`, `mesh`, and `voxel` collision geometry. The demo uses a cuboid. Create collision cache capacity before allowing world updates. +- Construct/cache a V2 `BatchMotionPlanner` with `max_batch_size == actual B`; do not reuse a larger graph-seeded batch planner for a smaller batch. +- `max_planning_time` is post-plan validation based on returned V2 timing; do not pass it to `MotionPlannerCfg.create`, which has no such argument. +- Do not linearly resample or interpolate a cuRobo collision-checked output after planning. Preserve the returned samples; action composition must accept their runtime length. + +--- + +## File Structure + +| File | Responsibility | +|---|---| +| `embodichain/lab/sim/planners/base_planner.py` | Add backend capabilities and generic motion-context hook. | +| `embodichain/lab/sim/planners/motion_generator.py` | Dispatch capability-aware interpolation and contextual options. | +| `embodichain/lab/sim/planners/neural_planner.py` | Migrate the existing Neural context behavior to the new hooks. | +| `embodichain/lab/sim/planners/toppra_planner.py` | Supply a correct default `ToppraPlanOptions`. | +| `embodichain/lab/sim/planners/curobo_planner.py` | Lazy V2 bindings, profile/world configs, named-joint conversion, scene updates, planning, and `PlanResult` conversion. | +| `embodichain/lab/sim/planners/__init__.py` | Re-export the CuRobo public API without importing cuRobo itself. | +| `embodichain/lab/sim/atomic_actions/core.py` | Validate `motion_source` / `planner_type` settings. | +| `embodichain/lab/sim/atomic_actions/trajectory.py` | Build CuRobo options, validate planner/result contracts, preserve collision-checked output, and add joint-motion dispatch. | +| `embodichain/lab/sim/atomic_actions/primitives/move_joints.py` | Opt in to collision-aware joint-space planning. | +| `embodichain/lab/sim/atomic_actions/primitives/{pick_up,place,press}.py` | Compose phase trajectories using actual returned lengths and avoid CuRobo pre-IK filtering. | +| `embodichain/lab/sim/atomic_actions/primitives/{coordinated_pickment,coordinated_placement}.py` | Fail explicitly if configured for the unsupported CuRobo coordinated path. | +| `embodichain/data/assets/curobo/collision_franka_demo.yml` | Static cuboid scene shared by the demo and optional integration test. | +| `tests/sim/planners/test_curobo_planner.py` | Dependency-free config, conversion, mapping, capability, and fake-backend tests. | +| `tests/sim/planners/test_curobo_integration.py` | CUDA/cuRobo V2 integration test, skipped when unavailable. | +| `tests/sim/atomic_actions/test_trajectory_motion_source.py` | CuRobo routing, no pre-IK, contract, and failure-hold tests. | +| `tests/sim/atomic_actions/test_actions.py` | `MoveJoints` CuRobo path and variable phase-length tests. | +| `tests/sim/atomic_actions/test_curobo_motion_source_e2e.py` | Optional DexSim endpoint test through `AtomicActionEngine`. | +| `examples/sim/planners/curobo_planner.py` | Runnable Panda + cuboid CuRobo V2 planner/action demo. | +| `docs/source/overview/sim/planners/{index.rst,curobo_planner.md,motion_generator.md}` | Public usage, installation, limitations, and planner overview. | +| `tests/docs/test_curobo_planner_docs.py` | Source-level documentation coverage check. | + +### Task 1: Generalize the Planner-to-MotionGenerator Protocol + +**Files:** +- Modify: `embodichain/lab/sim/planners/base_planner.py:31-195` +- Modify: `embodichain/lab/sim/planners/motion_generator.py:23-247` +- Modify: `embodichain/lab/sim/planners/neural_planner.py:194-350` +- Modify: `embodichain/lab/sim/planners/toppra_planner.py:237-399` +- Test: `tests/sim/planners/test_motion_generator_batched.py` +- Test: `tests/sim/planners/test_neural_planner.py` + +**Interfaces:** +- Consumes: existing `MotionGenOptions(start_qpos, control_part, plan_opts, is_interpolate)` and `PlanOptions`. +- Produces: `BasePlanner.preinterpolate_targets`, `BasePlanner.preserve_plan_samples`, `BasePlanner.default_plan_options()`, and `BasePlanner.with_motion_context(...)`; later `CuroboPlanner` implements these hooks. + +- [ ] **Step 1: Write failing capability-routing tests** + +Add a small fake planner to `tests/sim/planners/test_motion_generator_batched.py` and assert the original Cartesian target reaches a capability-false planner unchanged: + +```python +class _DirectCartesianPlanner: + preinterpolate_targets = False + preserve_plan_samples = True + + def default_plan_options(self): + return PlanOptions() + + def with_motion_context(self, options, *, start_qpos, control_part): + self.received = (start_qpos.clone(), control_part) + return options + + def plan(self, target_states, options): + self.target_states = target_states + return PlanResult( + success=torch.tensor([True]), + positions=torch.zeros(1, 3, 2), + ) + + +def test_direct_cartesian_planner_skips_preinterpolation(monkeypatch): + planner = _DirectCartesianPlanner() + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + start = torch.tensor([[0.1, -0.2]]) + goal = PlanState.from_xpos(torch.eye(4).unsqueeze(0)) + + result = generator.generate( + [goal], + MotionGenOptions( + start_qpos=start, + control_part="arm", + is_interpolate=True, + ), + ) + + assert result.success.item() + assert planner.target_states[0].move_type is MoveType.EEF_MOVE + assert torch.equal(planner.target_states[0].xpos, goal.xpos) + assert torch.equal(planner.received[0], start) + assert planner.received[1] == "arm" +``` + +Add a regression test in `test_neural_planner.py` that calls `MotionGenerator.generate()` with `MotionGenOptions.start_qpos` and `control_part` and verifies that the neural rollout starts at that qpos. + +- [ ] **Step 2: Run the new tests to verify the current implementation fails** + +Run: + +```bash +pytest tests/sim/planners/test_motion_generator_batched.py -q +pytest tests/sim/planners/test_neural_planner.py -q +``` + +Expected: the direct-planner test fails because `MotionGenerator` either performs interpolation or does not call the generic context hook. + +- [ ] **Step 3: Add the protocol hooks and migrate current planners** + +In `BasePlanner`, add the following complete default behavior directly below `__init__`: + +```python + preinterpolate_targets: bool = True + """Whether MotionGenerator may pre-interpolate targets for this backend.""" + + preserve_plan_samples: bool = False + """Whether callers must retain planner-returned sample points exactly.""" + + def default_plan_options(self) -> PlanOptions: + """Return backend-default planning options.""" + return PlanOptions() + + def with_motion_context( + self, + options: PlanOptions, + *, + start_qpos: torch.Tensor | None, + control_part: str | None, + ) -> PlanOptions: + """Attach MotionGenerator runtime context to backend options. + + The base planner has no context fields and therefore returns ``options`` + unchanged. Backends with contextual options override this method. + """ + return options +``` + +Implement the Neural override without mutating a caller's object unexpectedly: + +```python +class NeuralPlanner(BasePlanner): + preinterpolate_targets = False + + def with_motion_context( + self, + options: PlanOptions, + *, + start_qpos: torch.Tensor | None, + control_part: str | None, + ) -> NeuralPlanOptions: + if not isinstance(options, NeuralPlanOptions): + raise TypeError("NeuralPlanner requires NeuralPlanOptions") + if options.control_part is None: + options.control_part = control_part + if options.start_qpos is None: + options.start_qpos = start_qpos + return options +``` + +Add `ToppraPlanner.default_plan_options()` returning `ToppraPlanOptions()`. Remove the Neural-only propagation block from `MotionGenerator.generate()` and replace it with: + +```python + if options.is_interpolate and not self.planner.preinterpolate_targets: + logger.log_warning( + f"{type(self.planner).__name__} does not support MotionGenerator " + "pre-interpolation; disabling it." + ) + options.is_interpolate = False + + if options.plan_opts is None: + options.plan_opts = self.planner.default_plan_options() + options.plan_opts = self.planner.with_motion_context( + options.plan_opts, + start_qpos=options.start_qpos, + control_part=options.control_part, + ) +``` + +Keep the existing `EEF_MOVE`/`JOINT_MOVE` interpolation path unchanged when `preinterpolate_targets` is true. + +- [ ] **Step 4: Run focused tests and formatting** + +Run: + +```bash +black embodichain/lab/sim/planners/base_planner.py embodichain/lab/sim/planners/motion_generator.py embodichain/lab/sim/planners/neural_planner.py embodichain/lab/sim/planners/toppra_planner.py tests/sim/planners/test_motion_generator_batched.py tests/sim/planners/test_neural_planner.py +pytest tests/sim/planners/test_motion_generator_batched.py tests/sim/planners/test_neural_planner.py -q +``` + +Expected: PASS, with existing TOPPRA and Neural tests retaining their behavior. + +- [ ] **Step 5: Commit the protocol change** + +```bash +git add embodichain/lab/sim/planners/base_planner.py embodichain/lab/sim/planners/motion_generator.py embodichain/lab/sim/planners/neural_planner.py embodichain/lab/sim/planners/toppra_planner.py tests/sim/planners/test_motion_generator_batched.py tests/sim/planners/test_neural_planner.py +git commit -m "refactor(planner): add backend capability hooks" +``` + +### Task 2: Create the Optional CuRobo Configuration and Conversion Layer + +**Files:** +- Create: `embodichain/lab/sim/planners/curobo_planner.py` +- Modify: `embodichain/lab/sim/planners/__init__.py:17-21` +- Test: `tests/sim/planners/test_curobo_planner.py` + +**Interfaces:** +- Consumes: `BasePlannerCfg`, `PlanOptions`, `PlanState`, `PlanResult`, `configclass`, `quat_from_matrix`, and `SimulationManager` through `BasePlanner`. +- Produces: `CuroboRobotProfileCfg`, `CuroboWorldCfg`, `CuroboPlannerCfg`, `CuroboPlanOptions`, `CuroboPlanner`, plus dependency-free `_reorder_by_names`, `_matrix_to_position_quaternion`, and `_validate_dynamic_obstacles` helpers. + +- [ ] **Step 1: Write failing pure tests before adding the module** + +Create `tests/sim/planners/test_curobo_planner.py` with the project header and a fake `SimulationManager`. Add these representative cases: + +```python +def test_reorder_by_names_preserves_batch_and_time_dimensions(): + values = torch.tensor([[[10.0, 20.0], [30.0, 40.0]]]) + result = _reorder_by_names(values, ["joint_b", "joint_a"], ["joint_a", "joint_b"]) + assert torch.equal(result, torch.tensor([[[20.0, 10.0], [40.0, 30.0]]])) + + +def test_matrix_to_position_quaternion_uses_wxyz(): + matrix = torch.eye(4).unsqueeze(0) + position, quaternion = _matrix_to_position_quaternion(matrix) + assert torch.equal(position, torch.zeros(1, 3)) + assert torch.equal(quaternion, torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + + +def test_missing_curobo_is_actionable(monkeypatch): + monkeypatch.setattr(importlib, "import_module", _raise_module_not_found) + with pytest.raises(ImportError, match=r"cu12.*cu13"): + _require_curobo() + + +def test_unknown_dynamic_obstacle_is_rejected(): + with pytest.raises(ValueError, match="unknown obstacle"): + _validate_dynamic_obstacles({"unknown": torch.eye(4)}, ["known"]) +``` + +Also assert that `from embodichain.lab.sim.planners import CuroboPlannerCfg` succeeds without the installed `curobo` package. + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: + +```bash +pytest tests/sim/planners/test_curobo_planner.py -q +``` + +Expected: collection fails because `curobo_planner.py` and its exported types do not exist. + +- [ ] **Step 3: Implement lazy configs and pure helpers** + +Create `curobo_planner.py` with the project header, future import, `TYPE_CHECKING` guard, and this public surface: + +```python +__all__ = [ + "CuroboPlanOptions", + "CuroboPlanner", + "CuroboPlannerCfg", + "CuroboRobotProfileCfg", + "CuroboWorldCfg", +] + + +@configclass +class CuroboRobotProfileCfg: + robot_config_path: str = MISSING + sim_to_curobo_joint_names: dict[str, str] = MISSING + active_joint_names: list[str] | None = None + fixed_joint_positions: dict[str, float] = {} + base_link_name: str | None = None + tool_frame_name: str | None = None + + +@configclass +class CuroboWorldCfg: + world_config_path: str | None = None + collision_cache: dict[str, int] = {"cuboid": 8, "mesh": 2, "voxel": 1} + dynamic_obstacle_names: list[str] = [] + multi_env: bool = False + + +@configclass +class CuroboPlannerCfg(BasePlannerCfg): + planner_type: str = "curobo" + robot_profiles: dict[str, CuroboRobotProfileCfg] = MISSING + world: CuroboWorldCfg = CuroboWorldCfg() + warmup: bool = True + collision_activation_distance: float = 0.01 + max_attempts: int = 5 + max_planning_time: float | None = None + use_cuda_graph: bool = True + interpolation_dt: float = 0.025 + + +@configclass +class CuroboPlanOptions(PlanOptions): + start_qpos: torch.Tensor | None = None + control_part: str | None = None + dynamic_obstacle_poses: dict[str, torch.Tensor] | None = None + max_attempts: int | None = None +``` + +Implement `_require_curobo()` with `importlib.import_module` calls for the V2 public facades only. It must return an internal binding object containing `MotionPlanner`, `MotionPlannerCfg`, `BatchMotionPlanner`, `JointState`, `Pose`, and `GoalToolPose`; its `ImportError` must state both supported extras and the NVIDIA URL. Implement named reordering by validating equal unique name sets, not positional guesses. Implement matrix conversion with `embodichain.utils.math.quat_from_matrix`, whose output is wxyz, and reject non-`(B, 4, 4)` tensors. + +Append this lazy-safe export to `planners/__init__.py`: + +```python +from .curobo_planner import * +``` + +No statement at module scope may import `curobo`. + +- [ ] **Step 4: Pass the pure tests and check absence behavior** + +Run: + +```bash +black embodichain/lab/sim/planners/curobo_planner.py embodichain/lab/sim/planners/__init__.py tests/sim/planners/test_curobo_planner.py +python -c "from embodichain.lab.sim.planners import CuroboPlannerCfg; print(CuroboPlannerCfg.__name__)" +pytest tests/sim/planners/test_curobo_planner.py -q +``` + +Expected: imports print `CuroboPlannerCfg`; pure tests pass without cuRobo installed. + +- [ ] **Step 5: Commit the optional surface** + +```bash +git add embodichain/lab/sim/planners/curobo_planner.py embodichain/lab/sim/planners/__init__.py tests/sim/planners/test_curobo_planner.py +git commit -m "feat(planner): add optional curobo configuration" +``` + +### Task 3: Implement CuRobo V2 Backend Planning and World Synchronization + +**Files:** +- Modify: `embodichain/lab/sim/planners/curobo_planner.py` +- Modify: `tests/sim/planners/test_curobo_planner.py` +- Create: `tests/sim/planners/test_curobo_integration.py` + +**Interfaces:** +- Consumes: Task 1 planner hooks and Task 2 configs/helpers. +- Produces: `CuroboPlanner.plan(target_states, options) -> PlanResult`, `CuroboPlanner.update_dynamic_obstacles(poses)`, `CuroboPlanner.close()`, and a V2 backend cache keyed by `(control_part, batch_size, multi_env)`. + +- [ ] **Step 1: Write fake-binding tests for pose, c-space, output mapping, and failures** + +Extend `test_curobo_planner.py` with injected fake V2 bindings. Its fake result must model the high-level V2 shapes `[B, 1, T, D]` and named full-joint output: + +```python +def test_plan_pose_maps_curobo_full_output_to_control_part(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim) + result = planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.tensor([[0.2, -0.1]]), control_part="arm"), + ) + assert result.success.tolist() == [True] + assert result.positions.shape == (1, 3, 2) + assert torch.equal(result.positions[0, -1], torch.tensor([2.0, 1.0])) + assert result.dt.shape == (1, 3) + assert result.duration.shape == (1,) + + +def test_failed_v2_result_holds_start_qpos(fake_curobo, fake_sim): + fake_curobo.next_result.success = torch.tensor([[False]]) + planner = _make_planner(fake_curobo, fake_sim) + start = torch.tensor([[0.3, -0.4]]) + result = planner.plan([PlanState.from_qpos(start)], CuroboPlanOptions(start_qpos=start, control_part="arm")) + assert result.success.tolist() == [False] + assert torch.equal(result.positions, start.unsqueeze(1)) +``` + +Add a two-waypoint test that verifies the planner invokes V2 sequentially, begins the second segment at the first segment's final output, and concatenates the safe samples without a linear resample. Add tests for malformed V2 output names, unknown profile, unavailable CUDA, and a returned `total_time` greater than `max_planning_time` marking the affected result unsuccessful. + +- [ ] **Step 2: Run the fake-binding tests to verify they fail** + +Run: + +```bash +pytest tests/sim/planners/test_curobo_planner.py -q +``` + +Expected: fake V2 planning tests fail because `CuroboPlanner.plan()` is not implemented. + +- [ ] **Step 3: Implement backend construction and standard result conversion** + +Implement the following behavior in `CuroboPlanner`: + +```python +class CuroboPlanner(BasePlanner): + preinterpolate_targets = False + preserve_plan_samples = True + + def default_plan_options(self) -> CuroboPlanOptions: + return CuroboPlanOptions() + + def with_motion_context(self, options, *, start_qpos, control_part): + if not isinstance(options, CuroboPlanOptions): + raise TypeError("CuroboPlanner requires CuroboPlanOptions") + if options.start_qpos is None: + options.start_qpos = start_qpos + if options.control_part is None: + options.control_part = control_part + return options + + @validate_plan_options(options_cls=CuroboPlanOptions) + def plan(self, target_states, options=CuroboPlanOptions()) -> PlanResult: + control_part, profile = self._resolve_profile(options) + start = self._resolve_start_qpos(options.start_qpos, control_part) + backend = self._get_backend(profile, batch_size=start.shape[0]) + self.update_dynamic_obstacles(options.dynamic_obstacle_poses, backend) + return self._plan_segments(target_states, start, control_part, backend, options) +``` + +`_get_backend` must reject a non-CUDA robot device; invoke V2 `MotionPlannerCfg.create` with the profile path, `scene_model=world_config_path`, non-`None` `collision_cache`, `max_batch_size=batch_size`, `multi_env=cfg.world.multi_env`, `optimizer_collision_activation_distance`, `use_cuda_graph`, and `interpolation_dt`. Instantiate `MotionPlanner` for `B == 1` and `BatchMotionPlanner` for `B > 1`; call V2 warmup once per cache entry when `cfg.warmup` is true. Cache only entries matching the actual batch size. When `profile.active_joint_names` is set, compare it with `backend.planner.joint_names` and raise on any missing, duplicate, or differently ordered name. + +For each segment, construct V2 states and goals as follows: + +```python +active_start = _to_curobo_active_joint_state(start, profile, backend.planner.joint_names) +current_state = bindings.JointState.from_position( + active_start, + joint_names=backend.planner.joint_names, +) +position, quaternion = _matrix_to_position_quaternion(target.xpos) +goal = bindings.GoalToolPose.from_poses( + {backend.tool_frame: bindings.Pose(position=position, quaternion=quaternion)}, + ordered_tool_frames=[backend.tool_frame], + num_goalset=1, +) +v2_result = backend.planner.plan_pose(goal, current_state, max_attempts=max_attempts) +``` + +Use `plan_cspace(goal_state, current_state, ...)` for `JOINT_MOVE`. Map inputs by configured simulator-to-CuRobo joint names. V2 non-controlled joints must be locked in the robot profile or expanded/reduced through `backend.planner.kinematics.get_full_js` and `get_active_js`; never pass a nonexistent per-call fixed-joint argument. + +Extract `interpolated_trajectory` structurally, select seed zero from `[B, 1, T, D]`, trim each valid row using `interpolated_last_tstep`, map by `trajectory.joint_names` back to simulator control order, then concatenate sequential segment samples. To preserve a rectangular `PlanResult`, pad each batch row by repeating its own last valid qpos; set failed rows to `start_qpos.unsqueeze(1)`. Derive `dt` from V2 trajectory `dt` when it has time samples and otherwise from `cfg.interpolation_dt`; calculate `duration = dt.sum(dim=1)`. Do not expose private V2 `TrajOptSolverResult` in public annotations. + +`update_dynamic_obstacles` must validate the configured names and `(B, 4, 4)` shapes, convert to V2 `Pose` wxyz, and call `backend.planner.scene_collision_checker.update_obstacle_pose(name, pose, env_idx=index)` for each environment. `close()` must destroy every cached planner exactly once; `__del__` calls `close()` defensively. + +- [ ] **Step 4: Add and run the optional real V2 integration test** + +Create `tests/sim/planners/test_curobo_integration.py` with module-level guards: + +```python +curobo = pytest.importorskip("curobo") +if not torch.cuda.is_available(): + pytest.skip("cuRobo V2 requires CUDA", allow_module_level=True) + + +@pytest.mark.slow +def test_curobo_v2_plans_around_a_static_cuboid(...): + ... +``` + +The test must build a single Panda-compatible profile, call `MotionGenerator.generate()` for an EEF target around the cuboid, and assert `(B,)` success, finite `PlanResult` fields, correct first qpos, final FK position tolerance, output joint ordering, positive duration, and a reachable dynamic-obstacle update. Run: + +```bash +pytest tests/sim/planners/test_curobo_planner.py -q +pytest tests/sim/planners/test_curobo_integration.py -q +``` + +Expected: pure tests PASS; integration test is SKIPPED without V2/CUDA and PASS after the official V2 install. + +- [ ] **Step 5: Commit backend implementation and tests** + +```bash +git add embodichain/lab/sim/planners/curobo_planner.py tests/sim/planners/test_curobo_planner.py tests/sim/planners/test_curobo_integration.py +git commit -m "feat(planner): integrate curobo v2 backend" +``` + +### Task 4: Route Atomic Actions Through the Collision-Aware Backend + +**Files:** +- Modify: `embodichain/lab/sim/atomic_actions/core.py:294-315` +- Modify: `embodichain/lab/sim/atomic_actions/trajectory.py:293-475` +- Modify: `embodichain/lab/sim/atomic_actions/primitives/move_joints.py:78-133` +- Modify: `embodichain/lab/sim/atomic_actions/primitives/{pick_up,place,press}.py` +- Modify: `embodichain/lab/sim/atomic_actions/primitives/{coordinated_pickment,coordinated_placement}.py` +- Test: `tests/sim/atomic_actions/test_trajectory_motion_source.py` +- Test: `tests/sim/atomic_actions/test_actions.py` + +**Interfaces:** +- Consumes: `CuroboPlanOptions`, Task 1 capabilities, `ActionCfg`, and the existing `(success, controlled_trajectory)` builder contract. +- Produces: safe single-arm CuRobo atomic routing; `MoveJoints` motion-generator path; validation that prevents mislabeled planner use or a silent IK fallback. + +- [ ] **Step 1: Write failing atomic routing tests** + +Extend `test_trajectory_motion_source.py` with a fake CuRobo generator whose planner has `cfg.planner_type == "curobo"`, `preinterpolate_targets is False`, and `preserve_plan_samples is True`. Test all of these contracts: + +```python +def test_curobo_builder_preserves_cartesian_targets_and_samples(): + mg = _mock_curobo_motion_generator(result_positions=torch.zeros(2, 7, 6)) + builder = TrajectoryBuilder(mg) + success, trajectory = builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + n_waypoints=20, + control_part="arm", + arm_dof=6, + cfg=ActionCfg(motion_source="motion_gen", planner_type="curobo", control_part="arm"), + ) + assert success.tolist() == [True, True] + assert trajectory.shape == (2, 7, 6) + assert mg.generate.call_args.kwargs["options"].is_interpolate is False + assert mg.generate.call_args.args[0][0].move_type is MoveType.EEF_MOVE +``` + +Add tests that mismatched action/generator planner types, an invalid `motion_source`, malformed/NaN positions, and `None` positions raise `ValueError`. Add a failure-row test proving an unsuccessful row is held at its start qpos. + +In `test_actions.py`, add `MoveJointsCfg(motion_source="motion_gen", planner_type="curobo")` cases for a one-waypoint and multi-waypoint target. Assert ordered `JOINT_MOVE` PlanStates, full-DoF preservation of hand joints, returned CuRobo length, and per-env failure hold. + +- [ ] **Step 2: Run the atomic tests to verify they fail** + +Run: + +```bash +pytest tests/sim/atomic_actions/test_trajectory_motion_source.py tests/sim/atomic_actions/test_actions.py -q +``` + +Expected: CuRobo cases fail because the builder treats every non-Neural planner as pre-interpolated and `MoveJoints` always linearly interpolates. + +- [ ] **Step 3: Implement strict action and builder dispatch** + +Add `ActionCfg.__post_init__` validation: + +```python + def __post_init__(self) -> None: + valid_sources = {"ik_interp", "motion_gen"} + if self.motion_source not in valid_sources: + raise ValueError(f"motion_source must be one of {sorted(valid_sources)}") + if self.motion_source == "motion_gen" and self.planner_type is None: + raise ValueError("planner_type is required when motion_source='motion_gen'") + if self.motion_source == "ik_interp" and self.planner_type is not None: + raise ValueError("planner_type is only valid with motion_source='motion_gen'") +``` + +At the top of `TrajectoryBuilder._plan_motion_gen`, verify: + +```python +actual_type = self.motion_generator.planner.cfg.planner_type +requested_type = getattr(cfg, "planner_type", None) +if requested_type != actual_type: + raise ValueError( + f"Action requested planner_type={requested_type!r}, " + f"but MotionGenerator owns {actual_type!r}." + ) +``` + +Replace the `planner_type != "neural"` check with `not self.motion_generator.planner.preinterpolate_targets`. Replace `_build_plan_opts` with an explicit three-way factory: `ToppraPlanOptions` for TOPPRA, `NeuralPlanOptions` for Neural, and `CuroboPlanOptions(max_attempts=...)` for CuRobo. Reject unknown registered types rather than returning Neural options. + +Validate every motion-generator `PlanResult` before using it: + +```python +if positions is None or positions.ndim != 3: + raise ValueError("MotionGenerator returned no (B, N, controlled_dof) positions") +if positions.shape[0] != n_envs or positions.shape[2] != arm_dof: + raise ValueError("MotionGenerator returned incompatible trajectory shape") +if positions.device != self.device or not torch.isfinite(positions).all(): + raise ValueError("MotionGenerator returned non-finite or wrong-device positions") +``` + +If `planner.preserve_plan_samples` is true, return its trajectory unresampled; otherwise keep the existing `interpolate_with_distance` normalization. Add `TrajectoryBuilder.plan_joint_motion(...) -> tuple[torch.Tensor, torch.Tensor]`, which constructs batched `JOINT_MOVE` states and delegates to the same motion-generator validation path for `motion_gen`, or returns all-success linear interpolation for `ik_interp`. + +Modify `MoveJoints.execute` to call `plan_joint_motion` rather than directly calling `plan_joint_traj`: + +```python + success, joint_traj = self.builder.plan_joint_motion( + start_qpos, + target_qpos, + self.cfg.sample_interval, + control_part=self.cfg.control_part, + arm_dof=self.joint_dof, + cfg=self.cfg, + ) +``` + +Keep `_embed` and all non-controlled full qpos untouched. + +For `PickUp`, `Place`, and `Press`, allocate output phase tensors from `approach_arm.shape[1]`, `down_arm.shape[1]`, `back_arm.shape[1]`, and `lift_arm.shape[1]`, not their requested sample counts. Use `plan_joint_motion` for the `Press` return phase. When the configured planner is CuRobo, `PickUp._resolve_grasp_pose` selects the lowest-cost affordance variant without calling its existing EmbodiChain IK prefilter; CuRobo itself validates reachability/collision during motion planning. Raise immediately from coordinated primitive construction when the configuration requests `planner_type="curobo"`. + +- [ ] **Step 4: Run atomic unit tests and regression tests** + +Run: + +```bash +black embodichain/lab/sim/atomic_actions/core.py embodichain/lab/sim/atomic_actions/trajectory.py embodichain/lab/sim/atomic_actions/primitives/move_joints.py embodichain/lab/sim/atomic_actions/primitives/pick_up.py embodichain/lab/sim/atomic_actions/primitives/place.py embodichain/lab/sim/atomic_actions/primitives/press.py embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py tests/sim/atomic_actions/test_trajectory_motion_source.py tests/sim/atomic_actions/test_actions.py +pytest tests/sim/atomic_actions/test_trajectory_motion_source.py tests/sim/atomic_actions/test_actions.py tests/sim/atomic_actions/test_action_result_success.py -q +``` + +Expected: PASS; default `ik_interp` tests remain unchanged and CuRobo paths never invoke robot IK pre-interpolation. + +- [ ] **Step 5: Commit atomic integration** + +```bash +git add embodichain/lab/sim/atomic_actions/core.py embodichain/lab/sim/atomic_actions/trajectory.py embodichain/lab/sim/atomic_actions/primitives tests/sim/atomic_actions/test_trajectory_motion_source.py tests/sim/atomic_actions/test_actions.py +git commit -m "feat(actions): route atomic motions through curobo" +``` + +### Task 5: Add the Static-World Asset, Optional DexSim E2E Test, and Runnable Demo + +**Files:** +- Create: `embodichain/data/assets/curobo/collision_franka_demo.yml` +- Create: `tests/sim/atomic_actions/test_curobo_motion_source_e2e.py` +- Create: `examples/sim/planners/curobo_planner.py` + +**Interfaces:** +- Consumes: `FrankaPandaCfg`, `CuroboPlannerCfg`, `CuroboRobotProfileCfg`, `CuroboWorldCfg`, `MotionGenerator`, `AtomicActionEngine`, `MoveEndEffector`, and the Task 3 V2 backend. +- Produces: an executable one-environment cuboid-avoidance proof that exercises planner → motion generator → atomic action → DexSim playback. + +- [ ] **Step 1: Create the shared cuboid world asset and failing E2E test** + +Add `embodichain/data/assets/curobo/collision_franka_demo.yml`: + +```yaml +cuboid: + - name: demo_block + dims: [0.18, 0.40, 0.36] + pose: [0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0] +``` + +Create the e2e test with guards before any cuRobo-only imports: + +```python +curobo = pytest.importorskip("curobo") +if not torch.cuda.is_available(): + pytest.skip("cuRobo V2 requires CUDA", allow_module_level=True) + + +@pytest.mark.requires_sim +@pytest.mark.slow +def test_atomic_move_end_effector_uses_curobo_v2(): + sim, robot, engine = _make_franka_curobo_engine() + try: + target = _reachable_target_beyond_demo_block(robot) + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + assert success.shape == (1,) + assert success.item() + assert trajectory.shape[2] == robot.dof + _play_trajectory(sim, robot, trajectory) + assert _position_error(robot, target) < 0.02 + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() +``` + +- [ ] **Step 2: Run the E2E test to verify it fails before the fixture/demo exists** + +Run: + +```bash +pytest tests/sim/atomic_actions/test_curobo_motion_source_e2e.py -q +``` + +Expected: SKIPPED without cuRobo; after installation it initially fails because the shared engine fixture has not been written. + +- [ ] **Step 3: Implement the Panda profile fixture and CLI demo** + +Use `FrankaPandaCfg.from_dict({"uid": "curobo_franka"})`, not the legacy differently-named robot helper in `neural_planner.py`. Define the profile mapping explicitly so no index order is assumed: + +```python +FRANKA_SIM_TO_CUROBO = { + "fr3_joint1": "panda_joint1", + "fr3_joint2": "panda_joint2", + "fr3_joint3": "panda_joint3", + "fr3_joint4": "panda_joint4", + "fr3_joint5": "panda_joint5", + "fr3_joint6": "panda_joint6", + "fr3_joint7": "panda_joint7", +} + +profile = CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names=FRANKA_SIM_TO_CUROBO, + fixed_joint_positions={"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}, + base_link_name="panda_link0", + tool_frame_name="panda_hand", +) +``` + +In both fixture and example, instantiate a DexSim `CubeCfg(size=[0.18, 0.40, 0.36])` at `[0.45, 0.0, 0.18]` with the same UID-independent geometry as the YAML. Construct `MotionGenerator(MotionGenCfg(planner_cfg=CuroboPlannerCfg(...)))`, register `MoveEndEffector` configured with `motion_source="motion_gen"`, `planner_type="curobo"`, and `control_part="arm"`, then issue a target that requires passing the cuboid rather than entering it. + +The demo CLI must accept `--headless`, `--step-repeat`, `--hold-steps`, and `--no-warmup`; it must default to CUDA and raise a clear error when CUDA/cuRobo is absent. It prints success, trajectory shape, duration, and final Cartesian position error before replaying each full-DoF sample with `robot.set_qpos(qpos=trajectory[:, waypoint])` and `sim.update(step=step_repeat)`. + +- [ ] **Step 4: Run the optional real verification and example** + +After following NVIDIA's V2 installation command matching the local CUDA/PyTorch combination, run: + +```bash +pytest tests/sim/planners/test_curobo_integration.py tests/sim/atomic_actions/test_curobo_motion_source_e2e.py -q +python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1 +``` + +Expected: both tests PASS and the example prints a successful `(1, N, robot.dof)` trajectory with final position error below `0.02` m. + +- [ ] **Step 5: Commit the runnable proof** + +```bash +git add embodichain/data/assets/curobo/collision_franka_demo.yml tests/sim/atomic_actions/test_curobo_motion_source_e2e.py examples/sim/planners/curobo_planner.py +git commit -m "feat(example): add curobo planner simulation demo" +``` + +### Task 6: Document Installation, Profiles, and Supported Limits + +**Files:** +- Create: `docs/source/overview/sim/planners/curobo_planner.md` +- Modify: `docs/source/overview/sim/planners/index.rst` +- Modify: `docs/source/overview/sim/planners/motion_generator.md` +- Create: `tests/docs/test_curobo_planner_docs.py` + +**Interfaces:** +- Consumes: final public configs, demo command, and the V2 limitations from the approved design. +- Produces: discoverable end-user instructions that do not imply core dependency installation or unsupported attachment/dual-arm behavior. + +- [ ] **Step 1: Write the documentation assertions as a source-level test** + +Create `tests/docs/test_curobo_planner_docs.py` with the project header and a source-level check that the planner index includes `curobo_planner` and that the page contains all required literal API names: + +```python +def test_curobo_planner_docs_are_linked_and_scoped(): + index = Path("docs/source/overview/sim/planners/index.rst").read_text() + page = Path("docs/source/overview/sim/planners/curobo_planner.md").read_text() + assert "curobo_planner.md" in index + assert "CuroboPlannerCfg" in page + assert 'planner_type="curobo"' in page + assert "cuRobo V2" in page + assert "attached-object" in page +``` + +- [ ] **Step 2: Run the source-level test to verify it fails** + +Run: + +```bash +pytest tests/docs -q -k curobo +``` + +Expected: FAIL because the page and index entry do not exist. + +- [ ] **Step 3: Add concise user-facing documentation** + +Document all of the following in `curobo_planner.md`: + +```markdown +## Install cuRobo V2 + +Install cuRobo using NVIDIA's CUDA-matched extras, then verify it with +`python -c "import curobo; print(curobo.__version__)"`. EmbodiChain does not +install cuRobo as a core dependency. + +## Configure a control part + +Use `CuroboRobotProfileCfg.sim_to_curobo_joint_names` to map simulator names to +the V2 robot profile. Generate new collision-sphere/self-collision profiles +with V2 `RobotBuilder`; do not use a plain URDF alone. + +## Supported scope + +Single-arm `MoveEndEffector` and opt-in `MoveJoints` are supported. Static +cuboid, mesh, and voxel worlds are supported. Attached objects, automatic +scene extraction, coordinated dual-arm planning, ActionBank, and CPU execution +are not supported by this release. +``` + +Update the index to include the page and revise `motion_generator.md` so it describes cuRobo as an available collision-aware backend rather than a future capability. Include the exact demo command and a link to NVIDIA's official installation documentation. + +- [ ] **Step 4: Build docs and run all focused checks** + +Run: + +```bash +pytest tests/docs -q +cd docs && make html +``` + +Expected: PASS and `docs/build/html/index.html` contains the CuRobo planner page. + +- [ ] **Step 5: Commit documentation** + +```bash +git add docs/source/overview/sim/planners/curobo_planner.md docs/source/overview/sim/planners/index.rst docs/source/overview/sim/planners/motion_generator.md tests/docs/test_curobo_planner_docs.py +git commit -m "docs(planner): document curobo v2 integration" +``` + +### Task 7: Run the Whole-Change Quality Gate + +**Files:** +- Verify: every file created or modified by Tasks 1-6. + +**Interfaces:** +- Consumes: the complete implementation and existing planner/atomic regression suites. +- Produces: evidence that optional dependency behavior, default planner behavior, code style, and docs all remain valid. + +- [ ] **Step 1: Run static and targeted regression checks** + +Run: + +```bash +black --check --diff --color ./ +pytest tests/sim/planners/test_plan_state_batched.py tests/sim/planners/test_motion_generator_batched.py tests/sim/planners/test_neural_planner.py tests/sim/planners/test_curobo_planner.py -q +pytest tests/sim/atomic_actions/test_action_result_success.py tests/sim/atomic_actions/test_trajectory_motion_source.py tests/sim/atomic_actions/test_actions.py -q +python -c "import embodichain.lab.sim.planners" +``` + +Expected: PASS without cuRobo installed; V2-only modules remain lazily guarded. + +- [ ] **Step 2: Run V2/CUDA checks when the dependency is available** + +Run: + +```bash +python -c "import curobo; print(curobo.__version__)" +pytest tests/sim/planners/test_curobo_integration.py tests/sim/atomic_actions/test_curobo_motion_source_e2e.py -q +python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1 +``` + +Expected: PASS and the demo prints a successful full-DoF trajectory. If cuRobo is intentionally absent, record the two tests as skipped and do not claim the CUDA path passed. + +- [ ] **Step 3: Run the project pre-commit skill and inspect the final diff** + +Run the `/pre-commit-check` workflow, then: + +```bash +git diff --check +git status --short +git log --oneline --max-count=8 +``` + +Expected: no whitespace errors; only scoped CuRobo implementation, tests, example, docs, and this approved plan are present. diff --git a/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md b/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md new file mode 100644 index 00000000..17712a18 --- /dev/null +++ b/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md @@ -0,0 +1,282 @@ +# CuRobo V2 Motion-Planning Integration Design + +**Status:** Approved for implementation on 2026-07-11 + +## Goal + +Integrate NVIDIA cuRobo V2 as an optional, collision-aware motion-planning +backend for EmbodiChain. The integration must preserve EmbodiChain's existing +planner, motion-generator, and atomic-action contracts; provide a runnable +single-arm simulation demonstration; and remain importable and testable on +installations that do not have cuRobo or CUDA. + +## Context and Constraints + +EmbodiChain separates kinematics from planning: + +```text +RobotCfg / Robot + └── control parts and IK solvers +BasePlanner + └── MotionGenerator + └── AtomicActionEngine -> AtomicAction -> ActionResult +``` + +`BasePlanner.plan()` consumes env-batched `PlanState` waypoints and returns a +`PlanResult`. `MotionGenerator` owns the selected planner. Atomic actions use +`TrajectoryBuilder` to turn that result into a full-robot trajectory. The new +backend must fit this chain; it must not invoke `env.step()` or create an +alternative action-execution API. + +The integration targets **cuRobo V2 only**. IsaacLab's implementation is a +useful reference for backend boundaries, named-joint remapping, and collision +world lifecycle, but it uses the incompatible cuRobo V1 `MotionGen` API and is +not copied. + +cuRobo requires a CUDA-capable NVIDIA GPU. It remains an optional runtime +dependency because its CUDA extra must match the user's installed PyTorch and +driver. EmbodiChain will not add cuRobo to core dependencies or import it at +module import time. + +## Scope + +### Included + +- A cuRobo V2 planner backend registered as `planner_type="curobo"`. +- Cartesian pose planning and joint-space planning through the existing + `PlanState` / `PlanResult` interface. +- Single-arm control-part planning, including robust named-joint ordering, + fixed non-controlled joints, TCP, root-frame, and joint-limit handling. +- Batch-aware planning: the adapter accepts EmbodiChain's leading batch + dimension and uses cuRobo's V2 batch planner where multiple environments are + requested. The CUDA integration test establishes the single-environment + path; batch behavior is covered by adapter-level tests and runtime validation. +- Explicit static collision-world profiles and an explicit API to update named + dynamic obstacle poses. The first demo uses one static obstacle represented + in both DexSim and cuRobo. +- `MotionGenerator` propagation of planner-specific execution context + (`start_qpos` and `control_part`) without planner class special-cases. +- Atomic-action routing for single-arm Cartesian actions and opt-in + collision-aware `MoveJoints`. +- Unit, optional CUDA, and real-simulation end-to-end tests. +- A runnable demo at `examples/sim/planners/curobo_planner.py`. + +### Explicitly Excluded from This Change + +- cuRobo V1 compatibility or a hidden V1 fallback. +- Automatic conversion of every DexSim entity into a cuRobo collision object. +- Attached-object collision geometry and automatic attachment/detachment while + picking. The public context boundary is designed so this can be added later. +- Coordinated dual-arm planning and the current `CoordinatedPickment` path. +- The legacy Gym ActionBank system, which has a separate NumPy trajectory + representation and execution lifecycle. +- CPU fallback. A cuRobo planner construction request without CUDA fails with + an actionable error. + +## Architecture + +### Configuration and Public API + +The planner package adds focused, serializable config objects: + +```text +CuroboRobotProfileCfg + robot_config_path + active_joint_names + fixed_joint_positions + base_link_name + tool_frame_name + +CuroboWorldCfg + world_config_path + dynamic_obstacle_names + +CuroboPlannerCfg(BasePlannerCfg) + planner_type = "curobo" + robot_profiles: dict[str, CuroboRobotProfileCfg] + world: CuroboWorldCfg + warmup: bool + collision_activation_distance + max_attempts + max_planning_time + use_cuda_graph + +CuroboPlanOptions(PlanOptions) + start_qpos + control_part + dynamic_obstacle_poses: dict[str, Tensor[B, 4, 4]] + velocity_scale + acceleration_scale + max_attempts +``` + +Each `robot_profiles` key is an EmbodiChain control-part name. The selected +profile explicitly maps between simulator joint names and cuRobo joint names; +the backend never assumes that simulator indices and cuRobo indices are equal. +The profile also pins non-controlled joints, including gripper joints, to +current or configured values when the cuRobo robot model contains them. + +`CuroboPlanner` is a `BasePlanner` implementation. It lazily imports V2 +types, creates and warms a `MotionPlanner` for a one-environment request and a +V2 `BatchMotionPlanner` for a multi-environment request, then returns the +standard result fields: + +```text +success: bool tensor (B,) +positions: float tensor (B, N, controlled_dof) +velocities: float tensor (B, N, controlled_dof) or None +accelerations: float tensor (B, N, controlled_dof) or None +dt: float tensor (B, N) +duration: float tensor (B,) +xpos_list: float tensor (B, N, 4, 4) when FK is available +``` + +The planner supports `EEF_MOVE` by converting EmbodiChain world-frame matrices +to cuRobo target poses, including the selected robot base transform and TCP. +It supports `JOINT_MOVE` with cuRobo's joint-space planning API. Unsupported +move types fail before planning with a clear `ValueError`. + +### Motion-Generator Integration + +`MotionGenerator` registers `"curobo"` alongside existing planner types. +`BasePlanner` gains two explicit extension points: + +- `preinterpolate_targets: bool`, which defaults to `True`; +- `with_motion_context(options, *, start_qpos, control_part)`, whose base + implementation returns the supplied options unchanged. + +Planner capabilities replace the current `isinstance(NeuralPlanner)` decision: + +- planners declare whether EmbodiChain pre-interpolation is appropriate; +- CuRobo declares it is not appropriate for Cartesian targets, so it receives + the original `EEF_MOVE` targets and performs its own collision-aware IK and + trajectory optimization; +- planners receive the same runtime context through + `with_motion_context(...)`, rather than `MotionGenerator` copying fields + only for one concrete planner. + +This prevents a Cartesian action from being silently converted through +EmbodiChain IK before CuRobo sees it. It also ensures the start position and +the requested control part reach `CuroboPlanOptions` on every call. + +### Collision-World Boundary + +The initial collision world is explicit and deterministic. `CuroboWorldCfg` +points to a cuRobo V2 scene profile containing mesh and primitive obstacles. +At initialization the adapter loads this static scene once and warms the +planner. Before a plan, callers may supply poses keyed by the configured +dynamic obstacle names; the adapter verifies every name and updates only those +poses. Adding or removing geometry requires an explicit world reload rather +than a fragile implicit scene scan. + +All obstacle poses, goal poses, robot base poses, and tool poses use a single +documented world coordinate convention. The adapter converts EmbodiChain +4x4 matrices to cuRobo's position/quaternion representation at this boundary +and tests the quaternion ordering directly. + +### Atomic-Action Integration + +The existing data flow remains intact: + +```text +AtomicActionEngine + -> TrajectoryBuilder.plan_arm_traj + -> MotionGenerator.generate + -> CuroboPlanner.plan + -> PlanResult + -> full-DoF ActionResult trajectory +``` + +`ActionCfg.planner_type` documents and validates `"curobo"`. When the action +uses `motion_source="motion_gen"` and a CuRobo generator, the builder creates +`CuroboPlanOptions`, disables EmbodiChain Cartesian pre-interpolation, and +embeds the controlled joint plan into full robot DoF exactly as it does for +other planners. + +The supported first-release atomic surface is: + +- `MoveEndEffector`; +- movement phases of `PickUp`, `Place`, `Press`, and `MoveHeldObject` in a + static collision scene; +- `MoveJoints` when explicitly configured with `motion_source="motion_gen"`. + +Default action behavior remains unchanged (`motion_source="ik_interp"`). +Coordinated two-arm actions reject the CuRobo backend until they have a +separate multi-arm planning design. Because attached-object collision modeling +is excluded, users must not claim collision-aware carrying for a held object +until that later extension is implemented. + +### Optional Dependency Behavior + +The public planner package exports its configuration types without importing +cuRobo. Constructing or calling `CuroboPlanner` performs the dependency and +CUDA checks. A missing dependency raises an error that names the official V2 +installation command family (`.[cu12]` or `.[cu13]`) and links to the project +installation instructions. Tests that require the real library use +`pytest.importorskip("curobo")`. + +## Error Handling and Correctness Rules + +- Missing `robot_uid`, unknown control part, absent profile, joint-name + mismatch, invalid robot/world profile, or incompatible batch size fail before + a CUDA plan is launched. +- A failed cuRobo solution becomes a per-environment `success` tensor. Its + trajectory is held at `start_qpos`, matching the existing + `AtomicActionEngine` failed-environment behavior. +- Planner output is reordered from cuRobo named joints back to EmbodiChain + control-part order before it reaches `PlanResult`. +- Output tensors are moved to the robot's device and have the exact batch and + DoF shapes expected by `BasePlanner`. +- The implementation validates that an action's requested `planner_type` + matches the already constructed `MotionGenerator` planner; it never silently + plans with TOPPRA or NeuralPlanner under a CuRobo label. + +## Verification Strategy + +1. **Pure unit tests, no cuRobo required** + - config validation, registry/export behavior, optional-import errors; + - planner capability propagation through `MotionGenerator`; + - named-joint reorder and matrix/pose conversion helpers; + - `TrajectoryBuilder` and `MoveJoints` routing with a fake CuRobo backend. + +2. **Optional CUDA/cuRobo integration tests** + - create a V2 planner with a Panda profile and static obstacle; + - plan a collision-free pose trajectory and assert success, finite values, + start state, endpoint tolerance, joint ordering, and time fields; + - update a configured obstacle pose and verify the update reaches the + backend. + +3. **DexSim end-to-end test** + - create a single-arm robot and a matching static obstacle; + - execute `AtomicActionEngine + MoveEndEffector` with + `planner_type="curobo"`; + - assert a full-DoF trajectory, per-environment success, and target-pose + tolerance after playback. Mark this test `requires_sim` and `slow`. + +4. **Runnable demo** + - instantiate the same robot, obstacle, profile, planner, and atomic + action; + - plan, print success/shape/endpoint error, and replay the returned full-DoF + trajectory; + - fail early with installation or profile guidance when CuRobo/CUDA is not + available. + +## Documentation and Installation + +The integration documentation will identify CuRobo V2 as an optional CUDA +dependency and send users to NVIDIA's official installation flow. It will +explain how to generate a robot profile with `RobotBuilder`, how to configure +a static world, how to select `planner_type="curobo"`, and the first-release +limits on dynamic geometry, attachments, coordinated arms, and ActionBank. + +## Acceptance Criteria + +- `import embodichain.lab.sim.planners` works when cuRobo is absent. +- An installed V2/cuRobo CUDA environment plans a collision-free Panda path + through the EmbodiChain `MotionGenerator` API. +- `MoveEndEffector` reaches a target through `AtomicActionEngine` with a + full-DoF result trajectory and without EmbodiChain pre-IK. +- Existing TOPPRA, NeuralPlanner, default atomic actions, and ActionBank + behavior continue to pass their existing tests. +- `examples/sim/planners/curobo_planner.py` runs with documented V2 profiles + and demonstrates obstacle-aware planning and playback. From 023459623054a130d0128ac246780b973d4c1833 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 11:06:13 +0000 Subject: [PATCH 2/7] refactor(planner): add backend capability hooks Add BasePlanner.preinterpolate_targets / preserve_plan_samples / default_plan_options / with_motion_context so MotionGenerator is capability-driven rather than special-casing NeuralPlanner. Neural and TOPPRA migrate to the new hooks; behavior is preserved. Co-Authored-By: Claude --- embodichain/lab/sim/planners/base_planner.py | 44 +++++++++++++++ .../lab/sim/planners/motion_generator.py | 26 ++++----- .../lab/sim/planners/neural_planner.py | 19 +++++++ .../lab/sim/planners/toppra_planner.py | 4 ++ .../planners/test_motion_generator_batched.py | 54 +++++++++++++++++++ tests/sim/planners/test_neural_planner.py | 2 +- 6 files changed, 131 insertions(+), 18 deletions(-) diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index c71866f3..cda8f011 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -155,6 +155,50 @@ def __init__(self, cfg: BasePlannerCfg): self.device = self.robot.device + preinterpolate_targets: bool = True + """Whether ``MotionGenerator`` may pre-interpolate targets for this backend. + + Backends that perform their own collision-aware IK/trajectory optimization + (e.g. cuRobo) set this to ``False`` so the original Cartesian targets + reach ``plan`` unchanged rather than being converted through EmbodiChain IK. + """ + + preserve_plan_samples: bool = False + """Whether callers must retain this planner's returned sample points exactly. + + When ``True``, ``TrajectoryBuilder`` returns the planner's trajectory + without resampling, preserving collision-checked samples. When ``False`` + (the default), the builder may normalize the trajectory to a requested + waypoint count. + """ + + def default_plan_options(self) -> PlanOptions: + """Return backend-default planning options.""" + return PlanOptions() + + def with_motion_context( + self, + options: PlanOptions, + *, + start_qpos: torch.Tensor | None, + control_part: str | None, + ) -> PlanOptions: + """Attach MotionGenerator runtime context to backend options. + + The base planner has no context fields and therefore returns ``options`` + unchanged. Backends with contextual options override this method. + + Args: + options: The backend's planning options, already constructed (either + by the caller or via :meth:`default_plan_options`). + start_qpos: Optional starting joint configuration ``(B, DOF)``. + control_part: Optional control-part name. + + Returns: + The (possibly mutated) planning options carrying the context. + """ + return options + @validate_plan_options @abstractmethod def plan( diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index d8300495..6b8aa0a4 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -156,10 +156,10 @@ def generate( Returns: PlanResult containing the planned trajectory details. """ - if options.is_interpolate and isinstance(self.planner, NeuralPlanner): + if options.is_interpolate and not self.planner.preinterpolate_targets: logger.log_warning( - "is_interpolate=True is not supported with NeuralPlanner; " - "disabling interpolation." + f"{type(self.planner).__name__} does not support MotionGenerator " + "pre-interpolation; disabling it." ) options.is_interpolate = False @@ -227,20 +227,12 @@ def generate( target_plan_states = target_states if options.plan_opts is None: - if hasattr(self.planner, "default_plan_options"): - options.plan_opts = self.planner.default_plan_options() - else: - options.plan_opts = PlanOptions() - - # Propagate MotionGenOptions fields into NeuralPlanOptions so that callers - # can set control_part/start_qpos at the MotionGenerator level. - if isinstance(self.planner, NeuralPlanner) and isinstance( - options.plan_opts, NeuralPlanOptions - ): - if options.plan_opts.control_part is None: - options.plan_opts.control_part = options.control_part - if options.plan_opts.start_qpos is None: - options.plan_opts.start_qpos = options.start_qpos + options.plan_opts = self.planner.default_plan_options() + options.plan_opts = self.planner.with_motion_context( + options.plan_opts, + start_qpos=options.start_qpos, + control_part=options.control_part, + ) return self.planner.plan( target_states=target_plan_states, options=options.plan_opts diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index 90e917f9..9ebdb6cc 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -245,6 +245,9 @@ class NeuralPlanner(BasePlanner): KeyError: If the checkpoint is missing required keys. """ + preinterpolate_targets = False + """Neural rollouts consume raw EEF waypoints; pre-interpolation is disabled.""" + def __init__(self, cfg: NeuralPlannerCfg): super().__init__(cfg) @@ -256,6 +259,22 @@ def __init__(self, cfg: NeuralPlannerCfg): def default_plan_options(self) -> NeuralPlanOptions: return NeuralPlanOptions() + def with_motion_context( + self, + options: PlanOptions, + *, + start_qpos: torch.Tensor | None, + control_part: str | None, + ) -> NeuralPlanOptions: + """Forward MotionGenerator context into :class:`NeuralPlanOptions`.""" + if not isinstance(options, NeuralPlanOptions): + logger.log_error("NeuralPlanner requires NeuralPlanOptions", TypeError) + if options.control_part is None: + options.control_part = control_part + if options.start_qpos is None: + options.start_qpos = start_qpos + return options + def _load_checkpoint(self, checkpoint_path: Path) -> None: if not checkpoint_path.exists(): logger.log_error( diff --git a/embodichain/lab/sim/planners/toppra_planner.py b/embodichain/lab/sim/planners/toppra_planner.py index 5fc897f6..bafc5dda 100644 --- a/embodichain/lab/sim/planners/toppra_planner.py +++ b/embodichain/lab/sim/planners/toppra_planner.py @@ -305,6 +305,10 @@ def __init__(self, cfg: ToppraPlannerCfg): # by SimulationManager.destroy(), which skips every Python finalizer. # __del__ below only handles in-process GC of an abandoned planner. + def default_plan_options(self) -> ToppraPlanOptions: + """Return backend-default planning options.""" + return ToppraPlanOptions() + @staticmethod def _resolve_mp_context(mp_context: str | None, device: torch.device) -> str: """Return the multiprocessing start method to use for the worker pool. diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 4b950d21..9d4f6f06 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -24,9 +24,63 @@ MotionGenerator, MotionGenOptions, ) +from embodichain.lab.sim.planners.base_planner import PlanOptions from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType +class _DirectCartesianPlanner: + """Fake backend that consumes raw Cartesian targets (like cuRobo). + + Used to verify ``MotionGenerator`` skips pre-interpolation and forwards the + runtime context through the generic capability hooks rather than a + planner-class special case. + """ + + preinterpolate_targets = False + preserve_plan_samples = True + + def default_plan_options(self) -> PlanOptions: + return PlanOptions() + + def with_motion_context(self, options, *, start_qpos, control_part): + self.received = (start_qpos.clone(), control_part) + return options + + def plan(self, target_states, options): + self.target_states = target_states + return PlanResult( + success=torch.tensor([True]), + positions=torch.zeros(1, 3, 2), + ) + + +def test_direct_cartesian_planner_skips_preinterpolation(): + planner = _DirectCartesianPlanner() + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + start = torch.tensor([[0.1, -0.2]]) + goal = PlanState.from_xpos(torch.eye(4).unsqueeze(0)) + + result = generator.generate( + [goal], + MotionGenOptions( + start_qpos=start, + control_part="arm", + is_interpolate=True, + ), + ) + + assert result.success.item() + # The original EEF target reaches the planner unchanged - no IK, no + # pre-interpolation, no start-pose prepend. + assert planner.target_states[0].move_type is MoveType.EEF_MOVE + assert torch.equal(planner.target_states[0].xpos, goal.xpos) + # Runtime context is forwarded through the generic hook. + assert torch.equal(planner.received[0], start) + assert planner.received[1] == "arm" + + def _mock_planner(b=3, n=15, dofs=6): planner = Mock() planner.robot.num_instances = b diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index df14d754..298c0b39 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -345,7 +345,7 @@ def test_motion_generator_neural_auto_disables_interpolation( ) assert result.success.all().item() - assert "is_interpolate=True is not supported with NeuralPlanner" in caplog.text + assert "does not support MotionGenerator pre-interpolation" in caplog.text def test_safe_torch_load_roundtrip(tmp_path): From e14f16cda7e563e4d430ad15903aaf9956a7e14f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 11:11:13 +0000 Subject: [PATCH 3/7] feat(planner): add optional curobo configuration Add CuroboRobotProfileCfg/CuroboWorldCfg/CuroboPlannerCfg/CuroboPlanOptions and pure helpers (_reorder_by_names, _matrix_to_position_quaternion, _validate_dynamic_obstacles, _require_curobo). The package exports the configs without importing curobo; construction triggers the lazy V2 import. Co-Authored-By: Claude --- embodichain/lab/sim/planners/__init__.py | 1 + .../lab/sim/planners/curobo_planner.py | 377 ++++++++++++++++++ tests/sim/planners/test_curobo_planner.py | 134 +++++++ 3 files changed, 512 insertions(+) create mode 100644 embodichain/lab/sim/planners/curobo_planner.py create mode 100644 tests/sim/planners/test_curobo_planner.py diff --git a/embodichain/lab/sim/planners/__init__.py b/embodichain/lab/sim/planners/__init__.py index bff37bf7..c0782071 100644 --- a/embodichain/lab/sim/planners/__init__.py +++ b/embodichain/lab/sim/planners/__init__.py @@ -18,4 +18,5 @@ from .base_planner import * from .toppra_planner import * from .neural_planner import * +from .curobo_planner import * from .motion_generator import * diff --git a/embodichain/lab/sim/planners/curobo_planner.py b/embodichain/lab/sim/planners/curobo_planner.py new file mode 100644 index 00000000..e4fccd17 --- /dev/null +++ b/embodichain/lab/sim/planners/curobo_planner.py @@ -0,0 +1,377 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""Optional NVIDIA cuRobo V2 collision-aware motion-planning backend. + +This module is importable without cuRobo installed. Only constructing a +:class:`CuroboPlanner` triggers the lazy V2 import (and the actionable error +when cuRobo/CUDA are unavailable). cuRobo V2 is an optional runtime dependency; +EmbodiChain never imports it at module load time. + +The backend converts EmbodiChain's env-batched ``PlanState`` waypoints into +cuRobo V2 ``JointState`` / ``GoalToolPose`` calls, plans collision-aware +trajectories, and maps the result back into the standard ``PlanResult`` shape. +""" + +from __future__ import annotations + +import importlib +from dataclasses import MISSING +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import torch + +from embodichain.utils import configclass, logger +from embodichain.utils.math import quat_from_matrix + +from .base_planner import ( + BasePlanner, + BasePlannerCfg, + PlanOptions, + validate_plan_options, +) +from .utils import PlanResult, PlanState + +if TYPE_CHECKING: + from typing import Any + +__all__ = [ + "CuroboPlanOptions", + "CuroboPlanner", + "CuroboPlannerCfg", + "CuroboRobotProfileCfg", + "CuroboWorldCfg", +] + + +# cuRobo V2 installation extras documented at NVIDIA's installation page. +_CUROBO_INSTALL_URL = ( + "https://nvlabs.github.io/curobo/latest/getting-started/installation.html" +) + + +@configclass +class CuroboRobotProfileCfg: + """Per-control-part mapping between an EmbodiChain robot and a cuRobo model. + + Each ``CuroboPlannerCfg.robot_profiles`` key is an EmbodiChain control-part + name. The profile explicitly maps simulator joint names to cuRobo joint + names so the backend never assumes simulator and cuRobo joint indices agree. + Non-controlled joints (e.g. gripper joints) present in the cuRobo model are + pinned to ``fixed_joint_positions``. + """ + + robot_config_path: str = MISSING + """Path/identifier of the cuRobo V2 robot profile (e.g. ``"franka.yml"``).""" + + sim_to_curobo_joint_names: dict[str, str] = MISSING + """Mapping from simulator joint names to cuRobo joint names for this part.""" + + active_joint_names: list[str] | None = None + """Optional explicit cuRobo active-joint ordering, validated against the backend.""" + + fixed_joint_positions: dict[str, float] = {} + """cuRobo joint names pinned to a constant value (e.g. gripper finger joints).""" + + base_link_name: str | None = None + """cuRobo robot base link name.""" + + tool_frame_name: str | None = None + """cuRobo tool frame name used as the planning target.""" + + +@configclass +class CuroboWorldCfg: + """Static collision-world configuration for the cuRobo backend.""" + + world_config_path: str | None = None + """Path/identifier of a cuRobo V2 scene profile (cuboid/mesh/voxel obstacles).""" + + collision_cache: dict[str, int] = {"cuboid": 8, "mesh": 2, "voxel": 1} + """Per-geometry cache capacity created before world updates.""" + + dynamic_obstacle_names: list[str] = [] + """Obstacle names whose poses may be updated between plans.""" + + multi_env: bool = False + """Whether the cuRobo world is shared across multiple environments.""" + + +@configclass +class CuroboPlannerCfg(BasePlannerCfg): + """Configuration for the cuRobo V2 planner backend.""" + + planner_type: str = "curobo" + + robot_profiles: dict[str, CuroboRobotProfileCfg] = MISSING + """Control-part name -> profile mapping. The first release supports one part per request.""" + + world: CuroboWorldCfg = CuroboWorldCfg() + """Static collision-world configuration.""" + + warmup: bool = True + """Whether to warm each cached planner once at construction time.""" + + collision_activation_distance: float = 0.01 + """cuRobo collision activation distance (optimizer setting).""" + + max_attempts: int = 5 + """Default per-plan cuRobo attempt count.""" + + max_planning_time: float | None = None + """Post-plan validation budget (seconds). ``None`` skips the timing check.""" + + use_cuda_graph: bool = True + """Whether cuRobo may use CUDA graphs internally.""" + + interpolation_dt: float = 0.025 + """Interpolation step (seconds) used by cuRobo and as a dt fallback.""" + + +@configclass +class CuroboPlanOptions(PlanOptions): + """Per-plan options for :class:`CuroboPlanner`. + + ``start_qpos`` and ``control_part`` are populated from the + :class:`~embodichain.lab.sim.planners.motion_generator.MotionGenOptions` + runtime context via :meth:`CuroboPlanner.with_motion_context`. + """ + + start_qpos: torch.Tensor | None = None + """Planning start joint configuration ``(B, controlled_dof)``.""" + + control_part: str | None = None + """EmbodiChain control-part name to plan for.""" + + dynamic_obstacle_poses: dict[str, torch.Tensor] | None = None + """Per-obstacle world poses ``(B, 4, 4)`` keyed by configured name.""" + + max_attempts: int | None = None + """Per-plan override of ``CuroboPlannerCfg.max_attempts``.""" + + +# ============================================================================= +# Pure conversion / validation helpers (no cuRobo import required) +# ============================================================================= + + +def _reorder_by_names( + values: torch.Tensor, + from_names: list[str], + to_names: list[str], +) -> torch.Tensor: + """Reorder the trailing joint dimension of ``values`` by name. + + Args: + values: Tensor whose last dimension is ordered by ``from_names``. + from_names: Joint names describing the current trailing-axis order. + to_names: Desired joint-name order; must be a permutation of + ``from_names``. + + Returns: + Tensor with the trailing axis reordered to ``to_names``. + + Raises: + ValueError: If the two name sets are not equal as sets. + """ + if sorted(from_names) != sorted(to_names): + raise ValueError( + f"Cannot reorder joints: source names {from_names} and target names " + f"{to_names} are not the same set." + ) + if from_names == to_names: + return values + perm = [from_names.index(name) for name in to_names] + return values[..., perm] + + +def _matrix_to_position_quaternion( + matrix: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert a batched homogeneous pose to cuRobo ``(position, quaternion)``. + + Args: + matrix: Batched homogeneous transforms of shape ``(B, 4, 4)``. + + Returns: + Tuple of ``(position (B, 3), quaternion (B, 4))`` where the quaternion + is in cuRobo's ``(w, x, y, z)`` convention. + + Raises: + ValueError: If ``matrix`` is not a ``(B, 4, 4)`` tensor. + """ + if matrix.dim() != 3 or matrix.shape[-2:] != (4, 4): + raise ValueError( + f"Expected (B, 4, 4) pose matrices, got shape {tuple(matrix.shape)}." + ) + matrix = matrix.to(dtype=torch.float32) + position = matrix[:, :3, 3] + quaternion = quat_from_matrix(matrix[:, :3, :3]) # wxyz + return position, quaternion + + +def _validate_dynamic_obstacles( + poses: dict[str, torch.Tensor] | None, + allowed_names: list[str], +) -> None: + """Validate dynamic-obstacle pose names and shapes. + + Args: + poses: Mapping of obstacle name -> pose tensor. ``None`` is a no-op. + allowed_names: Obstacle names declared in :class:`CuroboWorldCfg`. + + Raises: + ValueError: If a name is not configured, or a pose is not ``(B, 4, 4)``. + """ + if poses is None: + return + for name, pose in poses.items(): + if name not in allowed_names: + raise ValueError( + f"unknown obstacle '{name}'; configured dynamic obstacles: " + f"{allowed_names}." + ) + if ( + not isinstance(pose, torch.Tensor) + or pose.dim() != 3 + or pose.shape[-2:] != (4, 4) + ): + got = tuple(pose.shape) if isinstance(pose, torch.Tensor) else type(pose) + raise ValueError( + f"dynamic obstacle '{name}' pose must be (B, 4, 4), got {got}." + ) + + +# ============================================================================= +# Lazy cuRobo V2 binding acquisition +# ============================================================================= + + +def _require_curobo() -> "Any": + """Lazily import and bundle the cuRobo V2 public facade types. + + Returns: + A namespace exposing ``MotionPlanner``, ``MotionPlannerCfg``, + ``BatchMotionPlanner``, ``JointState``, ``Pose``, and ``GoalToolPose``. + + Raises: + ImportError: If cuRobo V2 is not installed, with an actionable message + naming NVIDIA's CUDA-matched extras. + """ + try: + planner_mod = importlib.import_module("curobo.planner") + state_mod = importlib.import_module("curobo.types.state") + math_mod = importlib.import_module("curobo.types.math") + goal_mod = importlib.import_module("curobo.types.goal") + except ModuleNotFoundError as exc: + raise ImportError( + "cuRobo V2 is required for the 'curobo' planner but was not found. " + "Install it using NVIDIA's CUDA-matched extras, e.g. " + "`pip install .[cu12]` or `pip install .[cu13]` " + "(also `.[cu12-torch]` / `.[cu13-torch]`). " + f"See {_CUROBO_INSTALL_URL} for details." + ) from exc + return SimpleNamespace( + MotionPlanner=planner_mod.MotionPlanner, + MotionPlannerCfg=planner_mod.MotionPlannerCfg, + BatchMotionPlanner=planner_mod.BatchMotionPlanner, + JointState=state_mod.JointState, + Pose=math_mod.Pose, + GoalToolPose=goal_mod.GoalToolPose, + ) + + +# ============================================================================= +# CuroboPlanner +# ============================================================================= + + +class CuroboPlanner(BasePlanner): + r"""cuRobo V2 collision-aware motion-planning backend. + + The planner lazily imports cuRobo V2 at construction time, builds and caches + a V2 ``MotionPlanner`` (single-environment) or ``BatchMotionPlanner`` + (multi-environment) per ``(control_part, batch_size, multi_env)`` key, and + converts standard batched :class:`PlanState` inputs into V2 planning calls. + + Cartesian (``EEF_MOVE``) targets are forwarded to cuRobo unchanged - the + backend performs its own collision-aware IK and trajectory optimization, so + EmbodiChain pre-interpolation is disabled (``preinterpolate_targets=False``) + and returned collision-checked samples are preserved + (``preserve_plan_samples=True``). + + Args: + cfg: Configuration for the cuRobo planner. + + Raises: + ImportError: If cuRobo V2 is not installed. + ValueError: If ``robot_uid`` is missing or the robot is not found. + """ + + preinterpolate_targets = False + preserve_plan_samples = True + + def __init__(self, cfg: CuroboPlannerCfg) -> None: + super().__init__(cfg) + self.cfg: CuroboPlannerCfg = cfg + self._bindings = _require_curobo() + # Cached V2 backends keyed by (control_part, batch_size, multi_env). + self._backend_cache: dict[tuple, "Any"] = {} + + def default_plan_options(self) -> CuroboPlanOptions: + """Return backend-default planning options.""" + return CuroboPlanOptions() + + def with_motion_context( + self, + options: PlanOptions, + *, + start_qpos: torch.Tensor | None, + control_part: str | None, + ) -> CuroboPlanOptions: + """Forward MotionGenerator context into :class:`CuroboPlanOptions`.""" + if not isinstance(options, CuroboPlanOptions): + logger.log_error("CuroboPlanner requires CuroboPlanOptions", TypeError) + if options.start_qpos is None: + options.start_qpos = start_qpos + if options.control_part is None: + options.control_part = control_part + return options + + @validate_plan_options(options_cls=CuroboPlanOptions) + def plan( + self, + target_states: list[PlanState], + options: CuroboPlanOptions = CuroboPlanOptions(), + ) -> PlanResult: + r"""Plan a collision-aware trajectory through ``target_states``. + + .. note:: + Implemented in the backend-integration task. The method currently + raises ``NotImplementedError`` until the V2 planning path is wired. + + Args: + target_states: List of :class:`PlanState` waypoints. ``EEF_MOVE`` + entries carry ``xpos`` ``(B, 4, 4)``; ``JOINT_MOVE`` entries + carry ``qpos`` ``(B, controlled_dof)``. + options: :class:`CuroboPlanOptions` carrying the runtime context. + + Returns: + :class:`PlanResult` with env-batched tensors. Failed environments + hold ``start_qpos``. + """ + raise NotImplementedError("CuroboPlanner.plan is not implemented yet.") diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py new file mode 100644 index 00000000..a4b3b7f7 --- /dev/null +++ b/tests/sim/planners/test_curobo_planner.py @@ -0,0 +1,134 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""Dependency-free unit tests for the optional cuRobo planner surface. + +These tests never import the real ``curobo`` package. They cover config +validation, the public export behavior, the named-joint reorder helper, the +matrix -> position/quaternion conversion, the dynamic-obstacle validator, and +the actionable error raised when cuRobo is absent. +""" + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.lab.sim.planners.curobo_planner import ( + CuroboPlanOptions, + CuroboPlanner, + CuroboPlannerCfg as CuroboPlannerCfgDirect, + CuroboRobotProfileCfg, + CuroboWorldCfg, + _matrix_to_position_quaternion, + _require_curobo, + _reorder_by_names, + _validate_dynamic_obstacles, +) + + +def _raise_module_not_found(*args, **kwargs): + raise ModuleNotFoundError("curobo not installed") + + +def test_public_config_imports_without_curobo(): + """The planner package must export cuRobo configs without curobo installed.""" + assert CuroboPlannerCfg.__name__ == "CuroboPlannerCfg" + assert CuroboPlannerCfgDirect is CuroboPlannerCfg + assert CuroboPlannerCfg().planner_type == "curobo" + + +def test_reorder_by_names_preserves_batch_and_time_dimensions(): + values = torch.tensor([[[10.0, 20.0], [30.0, 40.0]]]) # (1, 2, 2) + result = _reorder_by_names(values, ["joint_b", "joint_a"], ["joint_a", "joint_b"]) + assert torch.equal(result, torch.tensor([[[20.0, 10.0], [40.0, 30.0]]])) + + +def test_reorder_by_names_rejects_mismatched_name_sets(): + values = torch.zeros(1, 2, 2) + with pytest.raises(ValueError, match="name"): + _reorder_by_names(values, ["joint_a", "joint_b"], ["joint_a", "joint_c"]) + + +def test_matrix_to_position_quaternion_uses_wxyz(): + matrix = torch.eye(4).unsqueeze(0) + position, quaternion = _matrix_to_position_quaternion(matrix) + assert torch.equal(position, torch.zeros(1, 3)) + assert torch.equal(quaternion, torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + + +def test_matrix_to_position_quaternion_rejects_non_4x4_batch(): + with pytest.raises(ValueError, match="4, 4"): + _matrix_to_position_quaternion(torch.zeros(3, 3)) + + +def test_missing_curobo_is_actionable(monkeypatch): + monkeypatch.setattr(importlib, "import_module", _raise_module_not_found) + with pytest.raises(ImportError, match=r"cu12.*cu13"): + _require_curobo() + + +def test_unknown_dynamic_obstacle_is_rejected(): + with pytest.raises(ValueError, match="unknown obstacle"): + _validate_dynamic_obstacles({"unknown": torch.eye(4)}, ["known"]) + + +def test_dynamic_obstacle_shape_is_validated(): + # (4, 4) is not batched -> rejected; the API requires (B, 4, 4). + with pytest.raises(ValueError, match="4, 4"): + _validate_dynamic_obstacles({"known": torch.eye(4)}, ["known"]) + + +def test_curobo_plan_options_carries_context_fields(): + opts = CuroboPlanOptions( + start_qpos=torch.zeros(2, 7), + control_part="arm", + max_attempts=3, + ) + assert opts.control_part == "arm" + assert opts.max_attempts == 3 + assert opts.start_qpos.shape == (2, 7) + + +def test_curobo_planner_cfg_defaults(): + cfg = CuroboPlannerCfg(robot_uid="franka") + assert cfg.planner_type == "curobo" + assert cfg.warmup is True + assert cfg.max_attempts == 5 + assert cfg.use_cuda_graph is True + assert isinstance(cfg.world, CuroboWorldCfg) + + +def test_curobo_robot_profile_cfg_requires_joint_map(): + cfg = CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={"a": "b"}, + ) + assert cfg.robot_config_path == "franka.yml" + assert cfg.sim_to_curobo_joint_names == {"a": "b"} + assert cfg.fixed_joint_positions == {} + + +def test_curobo_planner_class_is_lazy_import_safe(): + """Referencing the class must not import curobo.""" + import sys + + sys.modules.pop("curobo", None) + assert CuroboPlanner.__name__ == "CuroboPlanner" + assert "curobo" not in sys.modules From 3e478541bc13812bc06a74b4ca30a647ec30eabc Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 11:29:15 +0000 Subject: [PATCH 4/7] feat(planner): integrate curobo v2 backend Implement CuroboPlanner.plan (pose + cspace), backend cache keyed by (control_part, batch_size, multi_env), segment chaining without resampling, failure/over-budget hold, dynamic obstacle updates, and close(). Add fake-binding unit tests (require CUDA) and an optional real V2 integration test skipped without cuRobo/CUDA. Includes the shared cuboid world asset. Co-Authored-By: Claude --- .../assets/curobo/collision_franka_demo.yml | 11 + .../lab/sim/planners/curobo_planner.py | 465 +++++++++++++++++- tests/sim/planners/test_curobo_integration.py | 152 ++++++ tests/sim/planners/test_curobo_planner.py | 423 +++++++++++++++- 4 files changed, 1041 insertions(+), 10 deletions(-) create mode 100644 embodichain/data/assets/curobo/collision_franka_demo.yml create mode 100644 tests/sim/planners/test_curobo_integration.py diff --git a/embodichain/data/assets/curobo/collision_franka_demo.yml b/embodichain/data/assets/curobo/collision_franka_demo.yml new file mode 100644 index 00000000..3503d4b4 --- /dev/null +++ b/embodichain/data/assets/curobo/collision_franka_demo.yml @@ -0,0 +1,11 @@ +# cuRobo V2 static collision scene for the Franka Panda demo. +# +# A single cuboid obstacle placed in front of the robot. The same geometry is +# mirrored in DexSim (CubeCfg) by the demo and end-to-end test so that the +# planner's collision world and the simulator stay consistent. +# +# Pose convention: [x, y, z, qw, qx, qy, qz]. +cuboid: + - name: demo_block + dims: [0.18, 0.40, 0.36] + pose: [0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0] diff --git a/embodichain/lab/sim/planners/curobo_planner.py b/embodichain/lab/sim/planners/curobo_planner.py index e4fccd17..0e3f3874 100644 --- a/embodichain/lab/sim/planners/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo_planner.py @@ -29,7 +29,7 @@ from __future__ import annotations import importlib -from dataclasses import MISSING +from dataclasses import MISSING, dataclass from types import SimpleNamespace from typing import TYPE_CHECKING @@ -44,7 +44,7 @@ PlanOptions, validate_plan_options, ) -from .utils import PlanResult, PlanState +from .utils import MoveType, PlanResult, PlanState if TYPE_CHECKING: from typing import Any @@ -319,6 +319,7 @@ class CuroboPlanner(BasePlanner): Raises: ImportError: If cuRobo V2 is not installed. + RuntimeError: If the robot is not on a CUDA device. ValueError: If ``robot_uid`` is missing or the robot is not found. """ @@ -328,9 +329,15 @@ class CuroboPlanner(BasePlanner): def __init__(self, cfg: CuroboPlannerCfg) -> None: super().__init__(cfg) self.cfg: CuroboPlannerCfg = cfg + if self.device.type != "cuda": + raise RuntimeError( + "cuRobo V2 requires a CUDA device, but robot " + f"'{cfg.robot_uid}' is on {self.device}. Move the simulation " + "to a CUDA device before constructing the curobo planner." + ) self._bindings = _require_curobo() # Cached V2 backends keyed by (control_part, batch_size, multi_env). - self._backend_cache: dict[tuple, "Any"] = {} + self._backend_cache: dict[tuple, "_CuroboBackend"] = {} def default_plan_options(self) -> CuroboPlanOptions: """Return backend-default planning options.""" @@ -360,9 +367,11 @@ def plan( ) -> PlanResult: r"""Plan a collision-aware trajectory through ``target_states``. - .. note:: - Implemented in the backend-integration task. The method currently - raises ``NotImplementedError`` until the V2 planning path is wired. + ``EEF_MOVE`` waypoints are forwarded to cuRobo's ``plan_pose``; + ``JOINT_MOVE`` waypoints use ``plan_cspace``. Multi-waypoint plans + chain sequentially: each segment starts from the previous segment's + final sample, and the returned collision-checked samples are + concatenated without resampling. Args: target_states: List of :class:`PlanState` waypoints. ``EEF_MOVE`` @@ -371,7 +380,445 @@ def plan( options: :class:`CuroboPlanOptions` carrying the runtime context. Returns: - :class:`PlanResult` with env-batched tensors. Failed environments - hold ``start_qpos``. + :class:`PlanResult` with env-batched tensors. ``success`` is + ``(B,)`` bool; ``positions`` is ``(B, N, controlled_dof)``; + ``dt`` is ``(B, N)``; ``duration`` is ``(B,)``. Failed environments + (planning failure or ``total_time`` over budget) hold ``start_qpos``. """ - raise NotImplementedError("CuroboPlanner.plan is not implemented yet.") + if not target_states: + return PlanResult( + success=torch.zeros(0, dtype=torch.bool, device=self.device), + positions=None, + ) + control_part, profile = self._resolve_profile(options) + start = self._resolve_start_qpos(options.start_qpos, control_part) + backend = self._get_backend(profile, control_part, start.shape[0]) + self.update_dynamic_obstacles(options.dynamic_obstacle_poses, backend) + return self._plan_segments(target_states, start, backend, options) + + # ------------------------------------------------------------------ + # Profile / start resolution + # ------------------------------------------------------------------ + + def _resolve_profile( + self, options: CuroboPlanOptions + ) -> tuple[str, CuroboRobotProfileCfg]: + """Resolve the requested control part and its cuRobo profile.""" + control_part = options.control_part + if control_part is None: + logger.log_error("CuroboPlanOptions.control_part is required.", ValueError) + if control_part not in self.cfg.robot_profiles: + logger.log_error( + f"No cuRobo profile for control part '{control_part}'. " + f"Configured parts: {sorted(self.cfg.robot_profiles)}.", + ValueError, + ) + return control_part, self.cfg.robot_profiles[control_part] + + def _resolve_start_qpos( + self, start_qpos: torch.Tensor | None, control_part: str + ) -> torch.Tensor: + """Resolve the planning start qpos into ``(B, controlled_dof)``.""" + if start_qpos is None: + start_qpos = self.robot.get_qpos(name=control_part) + start_qpos = torch.as_tensor( + start_qpos, dtype=torch.float32, device=self.device + ) + if start_qpos.dim() == 1: + start_qpos = start_qpos.unsqueeze(0) + return start_qpos + + # ------------------------------------------------------------------ + # Backend construction / caching + # ------------------------------------------------------------------ + + def _get_backend( + self, + profile: CuroboRobotProfileCfg, + control_part: str, + batch_size: int, + ) -> "_CuroboBackend": + """Return a cached V2 backend for ``(control_part, batch_size, multi_env)``.""" + multi_env = self.cfg.world.multi_env + key = (control_part, int(batch_size), bool(multi_env)) + if key in self._backend_cache: + return self._backend_cache[key] + + world_cfg = self.cfg.world + collision_cache = ( + dict(world_cfg.collision_cache) if world_cfg.collision_cache else None + ) + planner_cfg = self._bindings.MotionPlannerCfg.create( + robot_config_path=profile.robot_config_path, + scene_model=world_cfg.world_config_path, + collision_cache=collision_cache, + max_batch_size=int(batch_size), + multi_env=bool(multi_env), + optimizer_collision_activation_distance=self.cfg.collision_activation_distance, + use_cuda_graph=bool(self.cfg.use_cuda_graph), + interpolation_dt=float(self.cfg.interpolation_dt), + ) + if batch_size == 1: + planner = self._bindings.MotionPlanner(planner_cfg) + else: + planner = self._bindings.BatchMotionPlanner( + planner_cfg, max_batch_size=int(batch_size) + ) + + if profile.active_joint_names is not None: + expected = list(profile.active_joint_names) + actual = list(planner.joint_names) + if expected != actual: + logger.log_error( + f"active_joint_names {expected} do not match cuRobo model " + f"joints {actual} (missing/duplicate/out-of-order).", + ValueError, + ) + + backend = _CuroboBackend( + planner=planner, + tool_frame=profile.tool_frame_name, + profile=profile, + batch_size=int(batch_size), + ) + if self.cfg.warmup: + planner.warmup() + self._backend_cache[key] = backend + return backend + + # ------------------------------------------------------------------ + # Segment planning + # ------------------------------------------------------------------ + + def _plan_segments( + self, + target_states: list[PlanState], + start: torch.Tensor, + backend: "_CuroboBackend", + options: CuroboPlanOptions, + ) -> PlanResult: + """Plan each waypoint segment sequentially and assemble a PlanResult.""" + B = start.shape[0] + D = start.shape[1] + max_attempts = ( + options.max_attempts + if options.max_attempts is not None + else self.cfg.max_attempts + ) + per_env_samples: list[list[torch.Tensor]] = [[] for _ in range(B)] + per_env_dt: list[list[torch.Tensor]] = [[] for _ in range(B)] + alive = torch.ones(B, dtype=torch.bool, device=self.device) + current = start.clone() + + for seg_idx, target in enumerate(target_states): + current_state = self._to_curobo_joint_state(current, backend) + if target.move_type == MoveType.EEF_MOVE: + if target.xpos is None: + logger.log_error( + f"Segment {seg_idx} EEF_MOVE target missing xpos.", + ValueError, + ) + goal = self._to_curobo_pose_goal(target.xpos, backend) + v2_result = backend.planner.plan_pose( + goal, current_state, max_attempts=max_attempts + ) + elif target.move_type == MoveType.JOINT_MOVE: + if target.qpos is None: + logger.log_error( + f"Segment {seg_idx} JOINT_MOVE target missing qpos.", + ValueError, + ) + goal_state = self._to_curobo_joint_goal(target.qpos, backend) + v2_result = backend.planner.plan_cspace( + goal_state, current_state, max_attempts=max_attempts + ) + else: + logger.log_error( + f"cuRobo does not support move_type {target.move_type}.", + ValueError, + ) + + seg_success, seg_positions, seg_dt = self._extract_segment( + v2_result, backend + ) + seg_success = seg_success.to(self.device) & alive + if self.cfg.max_planning_time is not None: + total_time = self._extract_total_time(v2_result, B) + over = total_time > float(self.cfg.max_planning_time) + seg_success = seg_success & (~over) + + for b in range(B): + if seg_idx == 0: + per_env_samples[b].append(seg_positions[b]) + per_env_dt[b].append(seg_dt[b]) + elif alive[b]: + # Drop the duplicate junction sample (== previous segment's + # final) so collision-checked samples are not duplicated. + per_env_samples[b].append(seg_positions[b, 1:]) + per_env_dt[b].append(seg_dt[b, 1:]) + else: + per_env_samples[b].append(seg_positions[b, -1:]) + per_env_dt[b].append(seg_dt[b, -1:]) + if seg_success[b]: + current[b] = seg_positions[b, -1] + alive = seg_success + + return self._assemble_result(per_env_samples, per_env_dt, start, alive, B, D) + + def _extract_segment( + self, v2_result: "Any", backend: "_CuroboBackend" + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Extract ``(success, positions, dt)`` for one V2 planning result. + + ``positions`` is ``(B, T, controlled_dof)`` in simulator control-part + order, trimmed to each env's last valid timestep and padded to a + rectangular batch by repeating the last valid sample. + """ + success = torch.as_tensor(v2_result.success) + if success.dim() == 2: + success = success.squeeze(-1) + success = success.to(torch.bool).to(self.device) + + traj = v2_result.interpolated_trajectory + position = torch.as_tensor(traj.position) + if position.dim() == 4: + position = position[:, 0, :, :] # select seed 0: (B, T, D_full) + + last_tstep = torch.as_tensor(v2_result.interpolated_last_tstep) + if last_tstep.dim() == 2: + last_tstep = last_tstep.squeeze(-1) + + B, T, _ = position.shape + max_len = max(int((last_tstep + 1).max().item()), 1) + full = torch.zeros( + B, max_len, position.shape[-1], device=self.device, dtype=torch.float32 + ) + for b in range(B): + length = min(int(last_tstep[b].item()) + 1, T, max_len) + full[b, :length] = position[b, :length].float().to(self.device) + if length < max_len: + full[b, length:] = position[b, length - 1].float().to(self.device) + + seg_positions = self._map_curobo_to_sim(full, traj.joint_names, backend.profile) + seg_dt = self._extract_dt(v2_result, traj, max_len, B) + return success, seg_positions, seg_dt + + def _map_curobo_to_sim( + self, + full_positions: torch.Tensor, + curobo_joint_names: list[str], + profile: CuroboRobotProfileCfg, + ) -> torch.Tensor: + """Map a full cuRobo trajectory to simulator control-part joint order.""" + sim_to_curobo = profile.sim_to_curobo_joint_names + cols: list[int] = [] + for sim_name in sim_to_curobo: + cu_name = sim_to_curobo[sim_name] + if cu_name not in curobo_joint_names: + logger.log_error( + f"cuRobo trajectory is missing active joint '{cu_name}' " + f"(mapped from sim joint '{sim_name}'); trajectory joints: " + f"{list(curobo_joint_names)}.", + ValueError, + ) + cols.append(curobo_joint_names.index(cu_name)) + return full_positions[..., cols].to(dtype=torch.float32) + + def _extract_dt( + self, v2_result: "Any", traj: "Any", max_len: int, B: int + ) -> torch.Tensor: + """Derive ``(B, max_len)`` per-sample dt from the V2 trajectory.""" + raw_dt = getattr(traj, "dt", None) + dt = None + if isinstance(raw_dt, torch.Tensor): + if raw_dt.dim() == 1: + dt = raw_dt.unsqueeze(0).expand(B, -1) + elif raw_dt.dim() == 2: + dt = raw_dt + if dt is None: + return torch.full( + (B, max_len), + float(self.cfg.interpolation_dt), + device=self.device, + dtype=torch.float32, + ) + T = dt.shape[-1] + out = torch.zeros(B, max_len, device=self.device, dtype=torch.float32) + length = min(T, max_len) + out[:, :length] = dt[:, :length].to(self.device) + return out + + def _extract_total_time(self, v2_result: "Any", B: int) -> torch.Tensor: + """Return a ``(B,)`` total planning time tensor for budget validation.""" + tt = v2_result.total_time + if isinstance(tt, torch.Tensor): + if tt.dim() == 0: + return tt.unsqueeze(0).expand(B).to(self.device) + if tt.dim() == 2: + tt = tt.squeeze(-1) + return tt[:B].to(self.device) + return torch.full((B,), float(tt), device=self.device) + + def _assemble_result( + self, + per_env_samples: list[list[torch.Tensor]], + per_env_dt: list[list[torch.Tensor]], + start: torch.Tensor, + alive: torch.Tensor, + B: int, + D: int, + ) -> PlanResult: + """Concatenate per-env segment samples into a rectangular PlanResult.""" + env_lengths: list[int] = [] + for b in range(B): + if alive[b]: + env_lengths.append(sum(s.shape[0] for s in per_env_samples[b])) + else: + env_lengths.append(1) + max_len = max(env_lengths) if env_lengths else 1 + + positions = torch.zeros(B, max_len, D, device=self.device, dtype=torch.float32) + dt = torch.zeros(B, max_len, device=self.device, dtype=torch.float32) + for b in range(B): + if alive[b]: + cat = torch.cat(per_env_samples[b], dim=0) + cat_dt = torch.cat(per_env_dt[b], dim=0) + length = cat.shape[0] + positions[b, :length] = cat + positions[b, length:] = cat[-1] + dt[b, : min(cat_dt.shape[0], max_len)] = cat_dt[:max_len] + else: + positions[b, :1] = start[b] + positions[b, 1:] = start[b] + duration = dt.sum(dim=1) + return PlanResult( + success=alive, + positions=positions, + dt=dt, + duration=duration, + ) + + # ------------------------------------------------------------------ + # cuRobo state / goal construction + # ------------------------------------------------------------------ + + def _to_curobo_joint_state( + self, current: torch.Tensor, backend: "_CuroboBackend" + ) -> "Any": + """Build a full cuRobo ``JointState`` from a sim-order control-part qpos. + + Active joints are filled from ``current`` (reordered to cuRobo order); + non-active joints present in the cuRobo model are pinned to + ``fixed_joint_positions``. + """ + profile = backend.profile + curobo_names = list(backend.planner.joint_names) + sim_to_curobo = profile.sim_to_curobo_joint_names + curobo_to_sim_idx = { + cu_name: idx + for idx, sim_name in enumerate(sim_to_curobo) + for cu_name in [sim_to_curobo[sim_name]] + } + B = current.shape[0] + state = torch.zeros( + B, len(curobo_names), device=self.device, dtype=torch.float32 + ) + for i, cu_name in enumerate(curobo_names): + if cu_name in curobo_to_sim_idx: + state[:, i] = current[:, curobo_to_sim_idx[cu_name]] + elif cu_name in profile.fixed_joint_positions: + state[:, i] = float(profile.fixed_joint_positions[cu_name]) + return self._bindings.JointState.from_position(state, joint_names=curobo_names) + + def _to_curobo_pose_goal( + self, xpos: torch.Tensor, backend: "_CuroboBackend" + ) -> "Any": + """Build a cuRobo ``GoalToolPose`` from a batched world-frame pose.""" + xpos = torch.as_tensor(xpos, device=self.device, dtype=torch.float32) + position, quaternion = _matrix_to_position_quaternion(xpos) + pose = self._bindings.Pose(position=position, quaternion=quaternion) + return self._bindings.GoalToolPose.from_poses( + {backend.tool_frame: pose}, + ordered_tool_frames=[backend.tool_frame], + num_goalset=1, + ) + + def _to_curobo_joint_goal( + self, qpos: torch.Tensor, backend: "_CuroboBackend" + ) -> "Any": + """Build a cuRobo c-space goal state from a sim-order target qpos.""" + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) + if qpos.dim() == 1: + qpos = qpos.unsqueeze(0) + full_state = self._to_curobo_joint_state(qpos, backend) + return self._bindings.JointState.from_position( + full_state, joint_names=list(backend.planner.joint_names) + ) + + # ------------------------------------------------------------------ + # Collision world + lifecycle + # ------------------------------------------------------------------ + + def update_dynamic_obstacles( + self, + poses: dict[str, torch.Tensor] | None, + backend: "_CuroboBackend | None" = None, + ) -> None: + """Update named dynamic obstacle poses on the cuRobo collision world. + + Args: + poses: Mapping of obstacle name -> ``(B, 4, 4)`` world pose. ``None`` + is a no-op. + backend: Specific backend to update. If ``None``, updates all cached + backends. + """ + if poses is None: + return + _validate_dynamic_obstacles(poses, list(self.cfg.world.dynamic_obstacle_names)) + backends = ( + [backend] if backend is not None else list(self._backend_cache.values()) + ) + for name, pose_tensor in poses.items(): + pose_tensor = torch.as_tensor( + pose_tensor, device=self.device, dtype=torch.float32 + ) + position, quaternion = _matrix_to_position_quaternion(pose_tensor) + B = pose_tensor.shape[0] + for b in range(B): + pose = self._bindings.Pose( + position=position[b], quaternion=quaternion[b] + ) + for be in backends: + be.planner.scene_collision_checker.update_obstacle_pose( + name, pose, env_idx=b + ) + + def close(self) -> None: + """Destroy every cached cuRobo planner and clear the cache.""" + for backend in list(self._backend_cache.values()): + planner = backend.planner + close_fn = getattr(planner, "close", None) or getattr( + planner, "destroy", None + ) + if close_fn is not None: + try: + close_fn() + except Exception: + pass + self._backend_cache.clear() + + def __del__(self) -> None: # pragma: no cover - best-effort GC cleanup + try: + self.close() + except Exception: + pass + + +@dataclass +class _CuroboBackend: + """Internal bundle of a cached V2 planner and its EmbodiChain-side metadata.""" + + planner: "Any" + tool_frame: str | None + profile: CuroboRobotProfileCfg + batch_size: int diff --git a/tests/sim/planners/test_curobo_integration.py b/tests/sim/planners/test_curobo_integration.py new file mode 100644 index 00000000..4ed2514f --- /dev/null +++ b/tests/sim/planners/test_curobo_integration.py @@ -0,0 +1,152 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""Optional cuRobo V2 + CUDA integration test. + +Skipped entirely when cuRobo or CUDA is unavailable. When both are present, +it builds a Panda profile + static cuboid world, plans a collision-aware EEF +move through the EmbodiChain ``MotionGenerator`` API, and verifies the +``PlanResult`` contract. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +# Module-level guards: skip the whole file without cuRobo or CUDA. These must +# run before any cuRobo-only import. +pytest.importorskip("curobo") +if not torch.cuda.is_available(): + pytest.skip("cuRobo V2 requires CUDA", allow_module_level=True) + +from embodichain import data as _data # noqa: E402 +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg # noqa: E402 +from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 +from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 +from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 +from embodichain.lab.sim.planners import ( # noqa: E402 + MotionGenCfg, + MotionGenOptions, + MotionGenerator, + MoveType, + PlanState, +) +from embodichain.lab.sim.planners.curobo_planner import ( # noqa: E402 + CuroboPlanOptions, + CuroboPlannerCfg, + CuroboRobotProfileCfg, + CuroboWorldCfg, +) + +ROBOT_UID = "curobo_franka" +CONTROL_PART = "arm" +DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] +DEMO_BLOCK_POS = [0.45, 0.0, 0.18] + + +def _demo_world_path() -> str: + return str( + Path(_data.__file__).parent / "assets" / "curobo" / "collision_franka_demo.yml" + ) + + +def _franka_profile() -> CuroboRobotProfileCfg: + """Explicit sim->cuRobo joint mapping; no index order is assumed.""" + sim_to_curobo = {f"fr3_joint{i}": f"panda_joint{i}" for i in range(1, 8)} + return CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names=sim_to_curobo, + fixed_joint_positions={ + "panda_finger_joint1": 0.04, + "panda_finger_joint2": 0.04, + }, + base_link_name="panda_link0", + tool_frame_name="panda_hand", + ) + + +def _make_sim_robot(): + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cuda", num_envs=1) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) + ) + # Mirror the cuRobo cuboid in DexSim so the planner and simulator agree. + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="demo_block", + shape=CubeCfg(size=DEMO_BLOCK_DIMS), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=DEMO_BLOCK_POS, + init_rot=[0.0, 0.0, 0.0], + ) + ) + return sim, robot + + +@pytest.mark.slow +def test_curobo_v2_plans_around_a_static_cuboid(): + sim, robot = _make_sim_robot() + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg(world_config_path=_demo_world_path()), + ) + mg = MotionGenerator(MotionGenCfg(planner_cfg=cfg)) + + start_qpos = robot.get_qpos(name=CONTROL_PART) + start_xpos = robot.compute_fk( + qpos=start_qpos, name=CONTROL_PART, to_matrix=True + ) + # Target beyond the cuboid so the planner must route around it. + target_xpos = start_xpos.clone() + target_xpos[0, :3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) + + result = mg.generate( + [PlanState.from_xpos(target_xpos)], + MotionGenOptions( + start_qpos=start_qpos, + control_part=CONTROL_PART, + plan_opts=CuroboPlanOptions(control_part=CONTROL_PART), + ), + ) + + assert result.success.shape == (1,) + assert bool(result.success.item()) + assert result.positions is not None + assert torch.isfinite(result.positions).all() + # Controlled joint count matches the arm (7). + assert result.positions.shape[-1] == 7 + # Trajectory starts at the requested start qpos. + assert torch.allclose(result.positions[0, 0], start_qpos[0], atol=1e-3) + # Positive duration. + assert float(result.duration[0]) > 0.0 + + # Final FK position reaches the target within tolerance. + final_q = result.positions[0, -1:].to(robot.device) + final_xpos = robot.compute_fk(qpos=final_q, name=CONTROL_PART, to_matrix=True) + err = torch.norm(final_xpos[0, :3, 3] - target_xpos[0, :3, 3]) + assert float(err) < 0.02 + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index a4b3b7f7..4dca2391 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -29,7 +29,7 @@ import pytest import torch -from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.lab.sim.planners import CuroboPlannerCfg, PlanState from embodichain.lab.sim.planners.curobo_planner import ( CuroboPlanOptions, CuroboPlanner, @@ -132,3 +132,424 @@ def test_curobo_planner_class_is_lazy_import_safe(): sys.modules.pop("curobo", None) assert CuroboPlanner.__name__ == "CuroboPlanner" assert "curobo" not in sys.modules + + +# ============================================================================= +# Fake cuRobo V2 bindings + backend planning tests +# ============================================================================= + +from embodichain.lab.sim.planners import curobo_planner as _curobo_mod +from embodichain.lab.sim.sim_manager import SimulationManager # noqa: E402 + + +class _FakeTrajectory: + def __init__(self, position, joint_names, dt=None): + self.position = position # (B, 1, T, D) + self.joint_names = list(joint_names) + self.dt = dt + + +class _FakeV2Result: + def __init__(self, success, trajectory, last_tstep, total_time=0.5): + self.success = success # (B, 1) + self.interpolated_trajectory = trajectory + self.interpolated_last_tstep = last_tstep # (B, 1) + self.total_time = total_time + + +class _FakeJointState: + def __init__(self, position, joint_names): + self.position = position + self.joint_names = joint_names + + +class _FakePose: + def __init__(self, position, quaternion): + self.position = position + self.quaternion = quaternion + + +class _FakeGoalToolPose: + def __init__(self, pose_dict, ordered_tool_frames): + self.pose_dict = pose_dict + self.ordered_tool_frames = ordered_tool_frames + + +class _FakeCollisionChecker: + def __init__(self): + self.updates = [] + + def update_obstacle_pose(self, name, pose, env_idx=0): + self.updates.append((name, pose, env_idx)) + + +class _FakeKinematics: + def __init__(self, joint_names): + self.joint_names = list(joint_names) + + +class _FakeV2PlannerInstance: + def __init__(self, bindings): + self._bindings = bindings + self.joint_names = list(bindings.full_joint_names) + self.tool_frame = bindings.tool_frame + self.scene_collision_checker = _FakeCollisionChecker() + self.kinematics = _FakeKinematics(self.joint_names) + self.plan_pose_calls = [] + self.plan_cspace_calls = [] + self.is_batch = False + self.max_batch_size = None + self.closed = False + self.warmup_count = 0 + + def plan_pose(self, goal, current_state, max_attempts=5): + self.plan_pose_calls.append((goal, current_state, max_attempts)) + return self._next_result() + + def plan_cspace(self, goal_state, current_state, max_attempts=5, **kwargs): + self.plan_cspace_calls.append((goal_state, current_state, max_attempts)) + return self._next_result() + + def _next_result(self): + if self._bindings.results: + return self._bindings.results.pop(0) + return self._bindings.next_result + + def warmup(self): + self.warmup_count += 1 + + def close(self): + self.closed = True + + +class _FakeMotionPlannerCfg: + def __init__(self, bindings): + self._bindings = bindings + + def create(self, **kwargs): + self._bindings.create_kwargs = kwargs + return ("fake_planner_cfg", kwargs) + + +class _FakeMotionPlanner: + def __init__(self, bindings): + self._bindings = bindings + + def __call__(self, cfg): + inst = _FakeV2PlannerInstance(self._bindings) + self._bindings.created_planners.append(inst) + return inst + + +class _FakeBatchMotionPlanner: + def __init__(self, bindings): + self._bindings = bindings + + def __call__(self, cfg, max_batch_size=None): + inst = _FakeV2PlannerInstance(self._bindings) + inst.is_batch = True + inst.max_batch_size = max_batch_size + self._bindings.created_planners.append(inst) + return inst + + +class _FakeJointStateFactory: + def __init__(self, bindings): + self._bindings = bindings + + def from_position(self, position, joint_names=None): + return _FakeJointState(position=position, joint_names=joint_names) + + +class _FakePoseFactory: + def __init__(self, bindings): + self._bindings = bindings + + def __call__(self, position, quaternion): + return _FakePose(position=position, quaternion=quaternion) + + +class _FakeGoalToolPoseFactory: + def __init__(self, bindings): + self._bindings = bindings + + def from_poses(self, pose_dict, ordered_tool_frames=None, num_goalset=1): + return _FakeGoalToolPose( + pose_dict=pose_dict, ordered_tool_frames=ordered_tool_frames + ) + + +class _FakeCuroboBindings: + """A minimal stand-in for the cuRobo V2 facade namespace.""" + + def __init__(self, full_joint_names, tool_frame="tool"): + self.full_joint_names = list(full_joint_names) + self.tool_frame = tool_frame + self.warmup_count = 0 + self.create_kwargs = None + self.created_planners: list = [] + self.results: list | None = None + self.MotionPlannerCfg = _FakeMotionPlannerCfg(self) + self.MotionPlanner = _FakeMotionPlanner(self) + self.BatchMotionPlanner = _FakeBatchMotionPlanner(self) + self.JointState = _FakeJointStateFactory(self) + self.Pose = _FakePoseFactory(self) + self.GoalToolPose = _FakeGoalToolPoseFactory(self) + self.next_result = self.make_result( + position=torch.zeros(1, 1, 3, len(full_joint_names)), + dt=torch.tensor([0.0, 0.025, 0.025]), + ) + + def make_result( + self, position, success=None, last_tstep=None, total_time=0.5, dt=None + ): + B, _, T, D = position.shape + if success is None: + success = torch.ones(B, 1, dtype=torch.bool) + if last_tstep is None: + last_tstep = torch.full((B, 1), T - 1, dtype=torch.long) + traj = _FakeTrajectory( + position=position, joint_names=list(self.full_joint_names), dt=dt + ) + return _FakeV2Result( + success=success, + trajectory=traj, + last_tstep=last_tstep, + total_time=total_time, + ) + + +class _FakeRobot: + def __init__(self, device="cuda", num_instances=1, dof=2): + self.uid = "fake_robot" + self.device = torch.device(device) + self.num_instances = num_instances + self.dof = dof + + def get_qpos(self, name=None): + return torch.zeros(self.num_instances, self.dof) + + def get_joint_ids(self, name=None): + return list(range(self.dof)) + + +class _FakeSim: + def __init__(self, robot): + self.robot = robot + + def get_robot(self, uid): + return self.robot + + +def _default_profile(): + return CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={"sim_a": "cu_a", "sim_b": "cu_b"}, + ) + + +def _make_planner( + fake_curobo, + fake_sim, + *, + profiles=None, + world=None, + **cfg_kw, +): + """Construct a CuroboPlanner against fake bindings + fake sim.""" + if profiles is None: + profiles = {"arm": _default_profile()} + cfg = CuroboPlannerCfg( + robot_uid="fake_robot", + robot_profiles=profiles, + world=world if world is not None else CuroboWorldCfg(), + **cfg_kw, + ) + return CuroboPlanner(cfg) + + +@pytest.fixture +def fake_sim(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("cuRobo backend requires a CUDA device") + robot = _FakeRobot(device="cuda") + sim = _FakeSim(robot) + monkeypatch.setattr(SimulationManager, "get_instance", classmethod(lambda cls: sim)) + return sim + + +@pytest.fixture +def fake_curobo(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("cuRobo backend requires a CUDA device") + bindings = _FakeCuroboBindings(full_joint_names=["cu_a", "cu_b"]) + monkeypatch.setattr(_curobo_mod, "_require_curobo", lambda: bindings) + return bindings + + +def test_plan_pose_maps_curobo_full_output_to_control_part(fake_curobo, fake_sim): + fake_curobo.next_result = fake_curobo.make_result( + position=torch.tensor([[[[0.2, -0.1], [1.5, 0.5], [2.0, 1.0]]]]), + dt=torch.tensor([0.0, 0.025, 0.025]), + ) + planner = _make_planner(fake_curobo, fake_sim) + result = planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.tensor([[0.2, -0.1]]), control_part="arm"), + ) + assert result.success.tolist() == [True] + assert result.positions.shape == (1, 3, 2) + assert torch.equal(result.positions[0, -1].cpu(), torch.tensor([2.0, 1.0])) + assert result.dt.shape == (1, 3) + assert result.duration.shape == (1,) + # warmup ran once and exactly one backend was built. + assert fake_curobo.created_planners[0].warmup_count == 1 + + +def test_failed_v2_result_holds_start_qpos(fake_curobo, fake_sim): + fake_curobo.next_result.success = torch.tensor([[False]]) + planner = _make_planner(fake_curobo, fake_sim) + start = torch.tensor([[0.3, -0.4]]) + result = planner.plan( + [PlanState.from_qpos(start)], + CuroboPlanOptions(start_qpos=start, control_part="arm"), + ) + assert result.success.tolist() == [False] + assert torch.equal(result.positions.cpu(), start.unsqueeze(1)) + + +def test_two_waypoints_chain_segments_without_resample(fake_curobo, fake_sim): + r1 = fake_curobo.make_result( + position=torch.tensor([[[[0.0, 0.0], [0.5, 0.0], [1.0, 0.0]]]]), + dt=torch.tensor([0.0, 0.025, 0.025]), + ) + r2 = fake_curobo.make_result( + position=torch.tensor([[[[1.0, 0.0], [2.0, 0.0]]]]), + dt=torch.tensor([0.0, 0.025]), + ) + fake_curobo.results = [r1, r2] + planner = _make_planner(fake_curobo, fake_sim) + start = torch.tensor([[0.0, 0.0]]) + result = planner.plan( + [ + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + ], + CuroboPlanOptions(start_qpos=start, control_part="arm"), + ) + assert result.success.tolist() == [True] + # seg1 (3) + seg2 without its duplicate junction (1) == 4 samples, unresampled. + assert result.positions.shape == (1, 4, 2) + assert torch.allclose(result.positions[0, 2].cpu(), torch.tensor([1.0, 0.0])) + assert torch.allclose(result.positions[0, -1].cpu(), torch.tensor([2.0, 0.0])) + # Segment 2 starts where segment 1 ended. + planner_inst = fake_curobo.created_planners[0] + assert len(planner_inst.plan_pose_calls) == 2 + second_current = planner_inst.plan_pose_calls[1][1].position + assert torch.allclose(second_current[0].cpu(), torch.tensor([1.0, 0.0])) + + +def test_malformed_trajectory_joint_names_raise(fake_curobo, fake_sim): + result = fake_curobo.make_result(position=torch.zeros(1, 1, 2, 2)) + result.interpolated_trajectory.joint_names = ["cu_a", "cu_x"] + fake_curobo.next_result = result + planner = _make_planner(fake_curobo, fake_sim) + with pytest.raises(ValueError, match="missing active joint"): + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + +def test_unknown_control_part_raises(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim) + with pytest.raises(ValueError, match="No cuRobo profile"): + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="leg"), + ) + + +def test_non_cuda_device_is_rejected(monkeypatch): + robot = _FakeRobot(device="cpu") + sim = _FakeSim(robot) + monkeypatch.setattr(SimulationManager, "get_instance", classmethod(lambda cls: sim)) + bindings = _FakeCuroboBindings(full_joint_names=["cu_a", "cu_b"]) + monkeypatch.setattr(_curobo_mod, "_require_curobo", lambda: bindings) + with pytest.raises(RuntimeError, match="CUDA"): + _make_planner(bindings, sim) + + +def test_total_time_over_budget_marks_unsuccessful(fake_curobo, fake_sim): + fake_curobo.next_result = fake_curobo.make_result( + position=torch.tensor([[[[0.0, 0.0], [1.0, 1.0], [2.0, 1.0]]]]), + total_time=0.5, + ) + planner = _make_planner(fake_curobo, fake_sim, max_planning_time=0.1) + start = torch.tensor([[0.3, -0.4]]) + result = planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=start, control_part="arm"), + ) + assert result.success.tolist() == [False] + assert torch.equal(result.positions.cpu(), start.unsqueeze(1)) + + +def test_backend_is_cached_across_plans(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim) + for _ in range(2): + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + assert len(fake_curobo.created_planners) == 1 + assert fake_curobo.created_planners[0].warmup_count == 1 + + +def test_update_dynamic_obstacles_reaches_backend(fake_curobo, fake_sim): + world = CuroboWorldCfg(dynamic_obstacle_names=["block"]) + planner = _make_planner(fake_curobo, fake_sim, world=world) + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + backend = next(iter(planner._backend_cache.values())) + planner.update_dynamic_obstacles( + {"block": torch.eye(4).unsqueeze(0).repeat(1, 1, 1)}, backend + ) + assert len(backend.planner.scene_collision_checker.updates) == 1 + name, _pose, env_idx = backend.planner.scene_collision_checker.updates[0] + assert name == "block" + assert env_idx == 0 + + +def test_close_destroys_cached_planners(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim) + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + assert len(fake_curobo.created_planners) == 1 + planner.close() + assert fake_curobo.created_planners[0].closed is True + assert planner._backend_cache == {} + + +def test_joint_move_uses_plan_cspace(fake_curobo, fake_sim): + fake_curobo.next_result = fake_curobo.make_result( + position=torch.tensor([[[[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]]]]), + dt=torch.tensor([0.0, 0.025, 0.025]), + ) + planner = _make_planner(fake_curobo, fake_sim) + start = torch.tensor([[0.0, 0.0]]) + target = torch.tensor([[1.0, 1.0]]) + result = planner.plan( + [PlanState.from_qpos(target)], + CuroboPlanOptions(start_qpos=start, control_part="arm"), + ) + assert result.success.tolist() == [True] + assert result.positions.shape == (1, 3, 2) + assert torch.allclose(result.positions[0, -1].cpu(), target[0]) + planner_inst = fake_curobo.created_planners[0] + assert len(planner_inst.plan_cspace_calls) == 1 + assert len(planner_inst.plan_pose_calls) == 0 From 8f4c96f196fab67c261341d0b77c2c502cad1bab Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 11:43:10 +0000 Subject: [PATCH 5/7] feat(actions): route atomic motions through curobo ActionCfg validates motion_source/planner_type. TrajectoryBuilder dispatches by capability (preinterpolate_targets / preserve_plan_samples) with strict planner-type matching, result validation, and a new plan_joint_motion path. MoveJoints, PickUp, Place, Press allocate from returned phase lengths; Press returns via plan_joint_motion; PickUp skips the IK prefilter for cuRobo. Coordinated dual-arm primitives reject the cuRobo backend. Co-Authored-By: Claude --- embodichain/lab/sim/atomic_actions/core.py | 20 ++- .../primitives/coordinated_pickment.py | 7 + .../primitives/coordinated_placement.py | 7 + .../atomic_actions/primitives/move_joints.py | 11 +- .../sim/atomic_actions/primitives/pick_up.py | 62 ++++++-- .../sim/atomic_actions/primitives/place.py | 20 ++- .../sim/atomic_actions/primitives/press.py | 25 ++- .../lab/sim/atomic_actions/trajectory.py | 132 ++++++++++++++-- tests/sim/atomic_actions/test_actions.py | 128 +++++++++++++++ .../test_trajectory_motion_source.py | 146 +++++++++++++++++- 10 files changed, 505 insertions(+), 53 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 4e80a4fa..b1191669 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -304,9 +304,27 @@ class ActionCfg: """Trajectory source: 'ik_interp' (default, batched IK + linear interp) or 'motion_gen' (batched MotionGenerator).""" planner_type: str | None = None - """Planner type for motion_source='motion_gen': 'toppra' | 'neural'. + """Planner type for motion_source='motion_gen': 'toppra' | 'neural' | 'curobo'. Required when motion_source='motion_gen'.""" + def __post_init__(self) -> None: + valid_sources = {"ik_interp", "motion_gen"} + if self.motion_source not in valid_sources: + raise ValueError( + f"motion_source must be one of {sorted(valid_sources)}, " + f"but got {self.motion_source!r}." + ) + if self.motion_source == "motion_gen" and self.planner_type is None: + raise ValueError( + "planner_type is required when motion_source='motion_gen'." + ) + if self.motion_source == "ik_interp" and self.planner_type is not None: + raise ValueError( + "planner_type is only valid with motion_source='motion_gen', " + f"but motion_source is 'ik_interp' and planner_type is " + f"{self.planner_type!r}." + ) + # ============================================================================= # AtomicAction ABC (slim) diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py index d8c2b9c7..2a43c1b8 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py @@ -410,6 +410,13 @@ def __init__( cfg: CoordinatedPickmentCfg | None = None, ) -> None: super().__init__(motion_generator, cfg or CoordinatedPickmentCfg()) + if getattr(self.cfg, "planner_type", None) == "curobo": + logger.log_error( + "Coordinated dual-arm planning is not supported by the cuRobo " + "backend. Use a single-arm action or a dedicated multi-arm " + "planner.", + ValueError, + ) self._init_dual_arm_parts( first_arm_control_part=self.cfg.left_arm_control_part, second_arm_control_part=self.cfg.right_arm_control_part, diff --git a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py index f8c2f88d..91eff826 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py +++ b/embodichain/lab/sim/atomic_actions/primitives/coordinated_placement.py @@ -102,6 +102,13 @@ def __init__( cfg: CoordinatedPlacementCfg | None = None, ) -> None: super().__init__(motion_generator, cfg or CoordinatedPlacementCfg()) + if getattr(self.cfg, "planner_type", None) == "curobo": + logger.log_error( + "Coordinated dual-arm planning is not supported by the cuRobo " + "backend. Use a single-arm action or a dedicated multi-arm " + "planner.", + ValueError, + ) self.builder = TrajectoryBuilder(motion_generator) self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index bf2b2968..845d6bb5 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -92,12 +92,17 @@ def execute( arm_dof=self.joint_dof, control_part=self.cfg.control_part, ) - joint_traj = self.builder.plan_joint_traj( - start_qpos, target_qpos, self.cfg.sample_interval + success, joint_traj = self.builder.plan_joint_motion( + start_qpos, + target_qpos, + self.cfg.sample_interval, + control_part=self.cfg.control_part, + arm_dof=self.joint_dof, + cfg=self.cfg, ) full = self._embed(joint_traj, state.last_qpos) return ActionResult( - success=torch.ones(self.n_envs, dtype=torch.bool, device=self.device), + success=success, trajectory=full, next_state=WorldState( last_qpos=full[:, -1, :].clone(), diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index bf290842..fb4a6e2d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -211,22 +211,29 @@ def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: self.hand_open_qpos, self.hand_close_qpos, n_waypoints=n_close ) + # Allocate from the actually-returned phase lengths so collision-aware + # planners (which preserve their own sample count) are not forced into + # the requested n_approach / n_lift counts. + n_approach_actual = approach_arm.shape[1] + n_lift_actual = lift_arm.shape[1] full = torch.empty( - (self.n_envs, n_approach + n_close + n_lift, self.robot_dof), + (self.n_envs, n_approach_actual + n_close + n_lift_actual, self.robot_dof), dtype=torch.float32, device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_approach, self.arm_joint_ids] = approach_arm - full[:, :n_approach, self.hand_joint_ids] = self.hand_open_qpos - full[:, n_approach : n_approach + n_close, self.arm_joint_ids] = ( + full[:, :n_approach_actual, self.arm_joint_ids] = approach_arm + full[:, :n_approach_actual, self.hand_joint_ids] = self.hand_open_qpos + full[:, n_approach_actual : n_approach_actual + n_close, self.arm_joint_ids] = ( grasp_arm_qpos.unsqueeze(1) ) - full[:, n_approach : n_approach + n_close, self.hand_joint_ids] = ( - hand_close_path + full[ + :, n_approach_actual : n_approach_actual + n_close, self.hand_joint_ids + ] = hand_close_path + full[:, n_approach_actual + n_close :, self.arm_joint_ids] = lift_arm + full[:, n_approach_actual + n_close :, self.hand_joint_ids] = ( + self.hand_close_qpos ) - full[:, n_approach + n_close :, self.arm_joint_ids] = lift_arm - full[:, n_approach + n_close :, self.hand_joint_ids] = self.hand_close_qpos obj_poses = sem.entity.get_local_pose(to_matrix=True) object_to_eef = torch.bmm(pose_inv(obj_poses), grasp_xpos) @@ -285,7 +292,7 @@ def _resolve_grasp_pose( grasp_xpos_padding[i, n_pose:] = grasp_poses[0] grasp_cost_padding[i, n_pose:] = grasp_costs[0] grasp_xpos_padding, ik_success = self._select_symmetric_grasp_variants( - grasp_xpos_padding, start_qpos + grasp_xpos_padding, start_qpos, skip_ik=self._is_motion_gen_curobo() ) grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) best_cost, best_idx = grasp_cost_masked.min(dim=1) @@ -295,10 +302,26 @@ def _resolve_grasp_pose( ] return is_success, best_grasp_xpos + def _is_motion_gen_curobo(self) -> bool: + """Whether this action is configured to plan through the cuRobo backend.""" + return ( + getattr(self.cfg, "motion_source", None) == "motion_gen" + and getattr(self.cfg, "planner_type", None) == "curobo" + ) + def _select_symmetric_grasp_variants( - self, grasp_xpos: torch.Tensor, start_qpos: torch.Tensor + self, + grasp_xpos: torch.Tensor, + start_qpos: torch.Tensor, + *, + skip_ik: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Choose the closest TCP z-roll variant, then validate reachability.""" + """Choose the closest TCP z-roll variant, then validate reachability. + + When ``skip_ik`` is set (cuRobo backend), the EmbodiChain IK prefilter + is skipped - cuRobo validates reachability and collision during motion + planning - and every variant is treated as reachable. + """ n_envs, n_pose = grasp_xpos.shape[:2] mirrored_grasp_xpos = grasp_xpos.clone() mirrored_grasp_xpos[..., :3, 0] = -mirrored_grasp_xpos[..., :3, 0] @@ -322,12 +345,17 @@ def _select_symmetric_grasp_variants( env_idx = torch.arange(n_envs, device=self.device)[:, None] pose_idx = torch.arange(n_pose, device=self.device)[None, :] selected_grasp_xpos = grasp_variants[env_idx, pose_idx, best_variant_idx] - start_qpos_repeat = start_qpos[:, None, :].repeat(1, n_pose, 1) - ik_success, _ = self.robot.compute_batch_ik( - pose=selected_grasp_xpos, - name=self.cfg.control_part, - joint_seed=start_qpos_repeat, - ) + if skip_ik: + ik_success = torch.ones( + n_envs, n_pose, dtype=torch.bool, device=self.device + ) + else: + start_qpos_repeat = start_qpos[:, None, :].repeat(1, n_pose, 1) + ik_success, _ = self.robot.compute_batch_ik( + pose=selected_grasp_xpos, + name=self.cfg.control_part, + joint_seed=start_qpos_repeat, + ) return selected_grasp_xpos, ik_success diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 942cdf3b..7e2fc58f 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -160,20 +160,26 @@ def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionRes self.hand_close_qpos, self.hand_open_qpos, n_waypoints=n_open ) + # Allocate from the actually-returned phase lengths so collision-aware + # planners (which preserve their own sample count) are accommodated. + n_down_actual = down_arm.shape[1] + n_back_actual = back_arm.shape[1] full = torch.empty( - (self.n_envs, n_down + n_open + n_back, self.robot_dof), + (self.n_envs, n_down_actual + n_open + n_back_actual, self.robot_dof), dtype=torch.float32, device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) - full[:, :n_down, self.arm_joint_ids] = down_arm - full[:, :n_down, self.hand_joint_ids] = self.hand_close_qpos - full[:, n_down : n_down + n_open, self.arm_joint_ids] = ( + full[:, :n_down_actual, self.arm_joint_ids] = down_arm + full[:, :n_down_actual, self.hand_joint_ids] = self.hand_close_qpos + full[:, n_down_actual : n_down_actual + n_open, self.arm_joint_ids] = ( reach_arm_qpos.unsqueeze(1) ) - full[:, n_down : n_down + n_open, self.hand_joint_ids] = hand_open_path - full[:, n_down + n_open :, self.arm_joint_ids] = back_arm - full[:, n_down + n_open :, self.hand_joint_ids] = self.hand_open_qpos + full[:, n_down_actual : n_down_actual + n_open, self.hand_joint_ids] = ( + hand_open_path + ) + full[:, n_down_actual + n_open :, self.arm_joint_ids] = back_arm + full[:, n_down_actual + n_open :, self.hand_joint_ids] = self.hand_open_qpos return ActionResult( success=success, diff --git a/embodichain/lab/sim/atomic_actions/primitives/press.py b/embodichain/lab/sim/atomic_actions/primitives/press.py index c30edbba..82d6fd44 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press.py @@ -115,23 +115,34 @@ def execute(self, target: EndEffectorPoseTarget, state: WorldState) -> ActionRes ) press_arm_qpos = down_arm[:, -1, :] - back_arm = self.builder.plan_joint_traj(press_arm_qpos, start_arm_qpos, n_back) - success = down_success + back_success, back_arm = self.builder.plan_joint_motion( + press_arm_qpos, + start_arm_qpos, + n_back, + control_part=self.cfg.control_part, + arm_dof=self.arm_dof, + cfg=self.cfg, + ) + success = down_success & back_success + # Allocate from the actually-returned phase lengths so collision-aware + # planners (which preserve their own sample count) are accommodated. + n_down_actual = down_arm.shape[1] + n_back_actual = back_arm.shape[1] full = torch.empty( - (self.n_envs, n_close + n_down + n_back, self.robot_dof), + (self.n_envs, n_close + n_down_actual + n_back_actual, self.robot_dof), dtype=torch.float32, device=self.device, ) full[:, :, :] = state.last_qpos.unsqueeze(1) full[:, :n_close, self.arm_joint_ids] = start_arm_qpos.unsqueeze(1) full[:, :n_close, self.hand_joint_ids] = hand_close_path - full[:, n_close : n_close + n_down, self.arm_joint_ids] = down_arm - full[:, n_close : n_close + n_down, self.hand_joint_ids] = ( + full[:, n_close : n_close + n_down_actual, self.arm_joint_ids] = down_arm + full[:, n_close : n_close + n_down_actual, self.hand_joint_ids] = ( self.hand_close_qpos.unsqueeze(1) ) - full[:, n_close + n_down :, self.arm_joint_ids] = back_arm - full[:, n_close + n_down :, self.hand_joint_ids] = ( + full[:, n_close + n_down_actual :, self.arm_joint_ids] = back_arm + full[:, n_close + n_down_actual :, self.hand_joint_ids] = ( self.hand_close_qpos.unsqueeze(1) ) diff --git a/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index e2fc7938..a6b217b0 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -380,38 +380,77 @@ def _plan_motion_gen( arm_dof: int, cfg: "ActionCfg | None", ) -> tuple[torch.Tensor, torch.Tensor]: - """Motion-generator trajectory source.""" + """Motion-generator trajectory source for Cartesian (EEF) targets.""" if self.motion_generator is None: logger.log_error( "motion_source='motion_gen' requires a MotionGenerator on the engine", ValueError, ) + self._validate_planner_type(cfg) n_envs = start_qpos.shape[0] plan_states = self._to_batched_plan_states(target_states_list, n_envs) plan_opts = self._build_plan_opts(cfg, n_waypoints) - planner_type = getattr(cfg, "planner_type", None) - is_interpolate = planner_type != "neural" result: PlanResult = self.motion_generator.generate( plan_states, - MotionGenOptions( + options=MotionGenOptions( start_qpos=start_qpos, control_part=control_part, plan_opts=plan_opts, - is_interpolate=is_interpolate, + is_interpolate=self.motion_generator.planner.preinterpolate_targets, ), ) + return self._process_motion_gen_result(result, start_qpos, n_waypoints, arm_dof) + + def _validate_planner_type(self, cfg: "ActionCfg | None") -> None: + """Reject actions whose requested planner differs from the engine's.""" + actual_type = self.motion_generator.planner.cfg.planner_type + requested_type = getattr(cfg, "planner_type", None) + if requested_type != actual_type: + logger.log_error( + f"Action requested planner_type={requested_type!r}, but " + f"MotionGenerator owns {actual_type!r}.", + ValueError, + ) + + def _process_motion_gen_result( + self, + result: PlanResult, + start_qpos: torch.Tensor, + n_waypoints: int, + arm_dof: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Validate a MotionGenerator PlanResult and apply sample/hold policy.""" success = ( result.success if isinstance(result.success, torch.Tensor) else torch.tensor(result.success, device=self.device) ) positions = result.positions - # Resample to n_waypoints if the planner returned a different count - if positions.shape[1] != n_waypoints: - positions = interpolate_with_distance( - trajectory=positions, interp_num=n_waypoints, device=self.device + n_envs = start_qpos.shape[0] + if positions is None or positions.ndim != 3: + logger.log_error( + "MotionGenerator returned no (B, N, controlled_dof) positions", + ValueError, + ) + if positions.shape[0] != n_envs or positions.shape[2] != arm_dof: + logger.log_error( + f"MotionGenerator returned incompatible trajectory shape " + f"{tuple(positions.shape)}; expected (..., {arm_dof}) on " + f"{n_envs} envs.", + ValueError, + ) + if positions.device != self.device or not torch.isfinite(positions).all(): + logger.log_error( + "MotionGenerator returned non-finite or wrong-device positions", + ValueError, ) - # Failed envs hold start qpos + if not self.motion_generator.planner.preserve_plan_samples: + if positions.shape[1] != n_waypoints: + positions = interpolate_with_distance( + trajectory=positions, interp_num=n_waypoints, device=self.device + ) + positions = positions.to(self.device) + # Failed envs hold start qpos across all waypoints. if not success.all(): held = start_qpos.unsqueeze(1).repeat(1, positions.shape[1], 1) positions = torch.where(success[:, None, None], positions, held) @@ -454,10 +493,10 @@ def _to_batched_plan_states( return batched def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): - """Build planner options from action configuration.""" + """Build planner options from action configuration (three-way factory).""" planner_type = getattr(cfg, "planner_type", None) - if planner_type in (None, "toppra"): - constraints = {} + if planner_type == "toppra": + constraints: dict = {} vl = getattr(cfg, "velocity_limit", None) al = getattr(cfg, "acceleration_limit", None) constraints["velocity"] = vl if vl is not None else 0.2 @@ -467,10 +506,71 @@ def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): sample_interval=n_waypoints, constraints=constraints, ) - # neural: planner reads its own cfg; pass minimal options - from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions + if planner_type == "neural": + from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions - return NeuralPlanOptions() + return NeuralPlanOptions() + if planner_type == "curobo": + from embodichain.lab.sim.planners.curobo_planner import CuroboPlanOptions + + return CuroboPlanOptions(max_attempts=getattr(cfg, "max_attempts", None)) + logger.log_error( + f"Unknown planner_type {planner_type!r} for motion_source='motion_gen'.", + ValueError, + ) + + def plan_joint_motion( + self, + start_qpos: torch.Tensor, + target_qpos: torch.Tensor, + n_waypoints: int, + *, + control_part: str, + arm_dof: int, + cfg: "ActionCfg | None" = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Plan a joint-space trajectory through one or more target waypoints. + + For ``motion_source='motion_gen'`` this delegates to the same + MotionGenerator validation path as Cartesian planning, building batched + ``JOINT_MOVE`` PlanStates. For ``motion_source='ik_interp'`` (the + default) it returns an all-success linear interpolation. + + Returns: + ``(success:(B,), trajectory:(B, N, arm_dof))``. + """ + motion_source = ( + getattr(cfg, "motion_source", "ik_interp") if cfg else "ik_interp" + ) + if motion_source == "motion_gen": + if self.motion_generator is None: + logger.log_error( + "motion_source='motion_gen' requires a MotionGenerator on the engine", + ValueError, + ) + self._validate_planner_type(cfg) + if target_qpos.dim() == 2: + target_qpos = target_qpos.unsqueeze(1) # (B, 1, D) + plan_states = [ + PlanState(qpos=target_qpos[:, j], move_type=MoveType.JOINT_MOVE) + for j in range(target_qpos.shape[1]) + ] + plan_opts = self._build_plan_opts(cfg, n_waypoints) + result: PlanResult = self.motion_generator.generate( + plan_states, + options=MotionGenOptions( + start_qpos=start_qpos, + control_part=control_part, + plan_opts=plan_opts, + is_interpolate=self.motion_generator.planner.preinterpolate_targets, + ), + ) + return self._process_motion_gen_result( + result, start_qpos, n_waypoints, arm_dof + ) + success = torch.ones(start_qpos.shape[0], dtype=torch.bool, device=self.device) + trajectory = self.plan_joint_traj(start_qpos, target_qpos, n_waypoints) + return success, trajectory def plan_joint_traj( self, diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 2ef969cd..1a88cbf6 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -25,6 +25,7 @@ from embodichain.lab.sim.atomic_actions.affordance import ( AntipodalAffordance, ) +from embodichain.lab.sim.planners.utils import MoveType, PlanResult from embodichain.lab.sim.atomic_actions.core import ( ActionResult, AtomicAction, @@ -123,6 +124,25 @@ def _make_mock_motion_generator(): return mg +def _make_curobo_mock_motion_generator(result_positions, success=None): + """Mock MotionGenerator whose planner is a cuRobo backend. + + ``result_positions`` is ``(B, N, ARM_DOF)``. The planner preserves samples + and disables pre-interpolation, matching the real cuRobo capabilities. + """ + mg = _make_mock_motion_generator() + planner = Mock() + planner.cfg.planner_type = "curobo" + planner.preinterpolate_targets = False + planner.preserve_plan_samples = True + mg.planner = planner + B = result_positions.shape[0] + if success is None: + success = torch.ones(B, dtype=torch.bool) + mg.generate.return_value = PlanResult(success=success, positions=result_positions) + return mg + + def _make_dual_arm_mock_robot(): robot = Mock() robot.device = torch.device("cpu") @@ -1038,3 +1058,111 @@ def interpolate(trajectory, interp_num, device): ) assert result.next_state.held_object.object_to_eef.shape == (NUM_ENVS, 4, 4) assert result.next_state.held_object.grasp_xpos.shape == (NUM_ENVS, 4, 4) + + +# --------------------------------------------------------------------------- +# MoveJoints + cuRobo motion_gen routing +# --------------------------------------------------------------------------- + + +class TestMoveJointsCurobo: + def setup_method(self): + # shared per-test mg is created in each test (result shapes differ). + pass + + def _action(self, mg, **cfg_kw): + return MoveJoints( + mg, + MoveJointsCfg(motion_source="motion_gen", planner_type="curobo", **cfg_kw), + ) + + def test_one_waypoint_routes_joint_move_to_motion_gen(self): + mg = _make_curobo_mock_motion_generator( + result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF) + ) + action = self._action(mg, sample_interval=10) + result = action.execute( + JointPositionTarget(qpos=torch.full((ARM_DOF,), 0.5)), + WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), + ) + assert result.success.tolist() == [True, True] + # Returned CuRobo length (5), not sample_interval (10). + assert result.trajectory.shape == (NUM_ENVS, 5, TOTAL_DOF) + # Full-DoF preservation: hand joints stay at the inherited state (zeros). + assert torch.allclose( + result.trajectory[:, :, ARM_DOF:], torch.zeros(NUM_ENVS, 5, HAND_DOF) + ) + plan_states = mg.generate.call_args.args[0] + assert all(s.move_type is MoveType.JOINT_MOVE for s in plan_states) + assert mg.generate.call_args.kwargs["options"].is_interpolate is False + + def test_multi_waypoint_routes_ordered_joint_states(self): + mg = _make_curobo_mock_motion_generator( + result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF) + ) + action = self._action(mg, sample_interval=10) + waypoint_qpos = ( + torch.stack( + [torch.full((ARM_DOF,), 0.3), torch.full((ARM_DOF,), 0.7)], dim=0 + ) + .unsqueeze(0) + .repeat(NUM_ENVS, 1, 1) + ) + result = action.execute( + JointPositionTarget(qpos=waypoint_qpos), + WorldState(last_qpos=torch.zeros(NUM_ENVS, TOTAL_DOF)), + ) + assert result.success.tolist() == [True, True] + assert result.trajectory.shape == (NUM_ENVS, 5, TOTAL_DOF) + plan_states = mg.generate.call_args.args[0] + assert len(plan_states) == 2 + assert all(s.move_type is MoveType.JOINT_MOVE for s in plan_states) + # Ordered: first waypoint, then second. + assert torch.allclose(plan_states[0].qpos, torch.full((NUM_ENVS, ARM_DOF), 0.3)) + assert torch.allclose(plan_states[1].qpos, torch.full((NUM_ENVS, ARM_DOF), 0.7)) + + def test_failure_holds_start_qpos(self): + positions = torch.zeros(NUM_ENVS, 5, ARM_DOF) + positions[1] = 1.0 # env 1 "would move" but is marked failed + mg = _make_curobo_mock_motion_generator( + result_positions=positions, success=torch.tensor([True, False]) + ) + action = self._action(mg, sample_interval=10) + last_qpos = torch.zeros(NUM_ENVS, TOTAL_DOF) + last_qpos[1, :ARM_DOF] = 0.7 # env 1 start + result = action.execute( + JointPositionTarget(qpos=torch.full((ARM_DOF,), 0.5)), + WorldState(last_qpos=last_qpos), + ) + assert result.success.tolist() == [True, False] + # Failed env held at its start arm qpos across all samples. + assert torch.allclose( + result.trajectory[1, :, :ARM_DOF], torch.full((5, ARM_DOF), 0.7) + ) + + +class TestCoordinatedRejectsCurobo: + def test_coordinated_pickment_rejects_curobo(self): + mg = _make_dual_arm_mock_motion_generator() + cfg = CoordinatedPickmentCfg( + left_hand_open_qpos=_hand_open(), + left_hand_close_qpos=_hand_close(), + right_hand_open_qpos=_hand_open(), + right_hand_close_qpos=_hand_close(), + motion_source="motion_gen", + planner_type="curobo", + ) + with pytest.raises(ValueError, match="not supported"): + CoordinatedPickment(mg, cfg) + + def test_coordinated_placement_rejects_curobo(self): + mg = _make_dual_arm_mock_motion_generator() + cfg = CoordinatedPlacementCfg( + placing_hand_open_qpos=_hand_open(), + placing_hand_close_qpos=_hand_close(), + support_hand_close_qpos=_hand_close(), + motion_source="motion_gen", + planner_type="curobo", + ) + with pytest.raises(ValueError, match="not supported"): + CoordinatedPlacement(mg, cfg) diff --git a/tests/sim/atomic_actions/test_trajectory_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py index cb3fccdc..dd552700 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -23,11 +23,11 @@ from unittest.mock import Mock from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder -from embodichain.lab.sim.planners.utils import PlanState, MoveType +from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType from embodichain.lab.sim.atomic_actions.core import ActionCfg -def _mock_mg(num_envs=2, arm_dof=6): +def _mock_mg(num_envs=2, arm_dof=6, planner_type="toppra"): robot = Mock() robot.device = torch.device("cpu") robot.dof = arm_dof @@ -40,9 +40,32 @@ def compute_ik(pose=None, name=None, joint_seed=None, **kw): mg = Mock() mg.robot = robot mg.device = torch.device("cpu") + planner = Mock() + planner.cfg.planner_type = planner_type + # TOPPRA allows MotionGenerator pre-interpolation; neural/curobo do not. + planner.preinterpolate_targets = planner_type == "toppra" + planner.preserve_plan_samples = planner_type == "curobo" + mg.planner = planner return mg +def _mock_curobo_motion_generator(result_positions, success=None): + """Fake MotionGenerator whose planner is a cuRobo backend.""" + num_envs = result_positions.shape[0] + arm_dof = result_positions.shape[-1] + mg = _mock_mg(num_envs=num_envs, arm_dof=arm_dof, planner_type="curobo") + if success is None: + success = torch.ones(num_envs, dtype=torch.bool) + mg.generate.return_value = PlanResult(success=success, positions=result_positions) + return mg + + +def _pose_targets_for_two_envs(): + return [ + [PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE)] for _ in range(2) + ] + + class TestPlanArmTrajMotionGen: def test_motion_gen_path_delegates_to_generate(self): mg = _mock_mg(num_envs=3, arm_dof=6) @@ -97,3 +120,122 @@ def test_ik_interp_path_unchanged(self): ) assert ok.all().item() assert traj.shape[0] == 2 + + +class TestCuroboBuilderDispatch: + def test_curobo_builder_preserves_cartesian_targets_and_samples(self): + mg = _mock_curobo_motion_generator(result_positions=torch.zeros(2, 7, 6)) + builder = TrajectoryBuilder(mg) + success, trajectory = builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + n_waypoints=20, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + assert success.tolist() == [True, True] + # preserve_plan_samples -> returned length is the planner's (7), not 20. + assert trajectory.shape == (2, 7, 6) + # No pre-interpolation; original EEF target reaches the generator. + assert mg.generate.call_args.kwargs["options"].is_interpolate is False + assert mg.generate.call_args.args[0][0].move_type is MoveType.EEF_MOVE + + def test_mismatched_planner_type_raises(self): + # MotionGenerator owns toppra, action requests curobo. + mg = _mock_mg(num_envs=2, arm_dof=6, planner_type="toppra") + builder = TrajectoryBuilder(mg) + with pytest.raises(ValueError, match="planner_type"): + builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + n_waypoints=10, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + + def test_invalid_motion_source_raises(self): + with pytest.raises(ValueError, match="motion_source"): + ActionCfg(motion_source="bogus") + + def test_motion_gen_without_planner_type_raises(self): + with pytest.raises(ValueError, match="planner_type is required"): + ActionCfg(motion_source="motion_gen") + + def test_ik_interp_with_planner_type_raises(self): + with pytest.raises(ValueError, match="planner_type is only valid"): + ActionCfg(motion_source="ik_interp", planner_type="toppra") + + def test_nan_positions_rejected(self): + positions = torch.zeros(2, 5, 6) + positions[0, 0, 0] = float("nan") + mg = _mock_curobo_motion_generator(result_positions=positions) + builder = TrajectoryBuilder(mg) + with pytest.raises(ValueError, match="non-finite"): + builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + n_waypoints=10, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + + def test_none_positions_rejected(self): + mg = _mock_curobo_motion_generator(result_positions=torch.zeros(2, 5, 6)) + mg.generate.return_value = PlanResult( + success=torch.ones(2, dtype=torch.bool), positions=None + ) + builder = TrajectoryBuilder(mg) + with pytest.raises(ValueError, match="positions"): + builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6), + n_waypoints=10, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + + def test_failed_row_holds_start_qpos(self): + positions = torch.zeros(2, 5, 6) + positions[1] = 1.0 # env 1 "succeeds" numerically but we mark it failed + mg = _mock_curobo_motion_generator( + result_positions=positions, + success=torch.tensor([True, False]), + ) + builder = TrajectoryBuilder(mg) + start = torch.zeros(2, 6) + start[1] = 0.5 + success, trajectory = builder.plan_arm_traj( + _pose_targets_for_two_envs(), + start, + n_waypoints=10, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + assert success.tolist() == [True, False] + # Failed env held at its start qpos across all samples. + assert torch.allclose(trajectory[1], start[1].unsqueeze(0).repeat(5, 1)) From 62f9edc83a64f399297df3640093a9bfacc05d9e Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 12:07:11 +0000 Subject: [PATCH 6/7] wip --- .../lab/sim/planners/curobo_planner.py | 21 +- .../lab/sim/planners/motion_generator.py | 4 + examples/sim/motion_gen/curobo_motion_gen.py | 225 ++++++++++++++++++ .../test_curobo_motion_source_e2e.py | 169 +++++++++++++ 4 files changed, 407 insertions(+), 12 deletions(-) create mode 100644 examples/sim/motion_gen/curobo_motion_gen.py create mode 100644 tests/sim/atomic_actions/test_curobo_motion_source_e2e.py diff --git a/embodichain/lab/sim/planners/curobo_planner.py b/embodichain/lab/sim/planners/curobo_planner.py index 0e3f3874..bb404e65 100644 --- a/embodichain/lab/sim/planners/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo_planner.py @@ -273,10 +273,9 @@ def _require_curobo() -> "Any": naming NVIDIA's CUDA-matched extras. """ try: - planner_mod = importlib.import_module("curobo.planner") - state_mod = importlib.import_module("curobo.types.state") - math_mod = importlib.import_module("curobo.types.math") - goal_mod = importlib.import_module("curobo.types.goal") + planner_mod = importlib.import_module("curobo.motion_planner") + batch_mod = importlib.import_module("curobo.batch_motion_planner") + types_mod = importlib.import_module("curobo.types") except ModuleNotFoundError as exc: raise ImportError( "cuRobo V2 is required for the 'curobo' planner but was not found. " @@ -288,10 +287,10 @@ def _require_curobo() -> "Any": return SimpleNamespace( MotionPlanner=planner_mod.MotionPlanner, MotionPlannerCfg=planner_mod.MotionPlannerCfg, - BatchMotionPlanner=planner_mod.BatchMotionPlanner, - JointState=state_mod.JointState, - Pose=math_mod.Pose, - GoalToolPose=goal_mod.GoalToolPose, + BatchMotionPlanner=batch_mod.BatchMotionPlanner, + JointState=types_mod.JointState, + Pose=types_mod.Pose, + GoalToolPose=types_mod.GoalToolPose, ) @@ -449,7 +448,7 @@ def _get_backend( dict(world_cfg.collision_cache) if world_cfg.collision_cache else None ) planner_cfg = self._bindings.MotionPlannerCfg.create( - robot_config_path=profile.robot_config_path, + robot=profile.robot_config_path, scene_model=world_cfg.world_config_path, collision_cache=collision_cache, max_batch_size=int(batch_size), @@ -461,9 +460,7 @@ def _get_backend( if batch_size == 1: planner = self._bindings.MotionPlanner(planner_cfg) else: - planner = self._bindings.BatchMotionPlanner( - planner_cfg, max_batch_size=int(batch_size) - ) + planner = self._bindings.BatchMotionPlanner(planner_cfg) if profile.active_joint_names is not None: expected = list(profile.active_joint_names) diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index 6b8aa0a4..f9446cff 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -30,6 +30,9 @@ NeuralPlanner, NeuralPlannerCfg, NeuralPlanOptions, + CuroboPlanner, + CuroboPlannerCfg, + CuroboPlanOptions, ) from embodichain.lab.sim.utility.action_utils import interpolate_with_nums from embodichain.utils import logger, configclass @@ -101,6 +104,7 @@ class MotionGenerator: _support_planner_dict = { "toppra": (ToppraPlanner, ToppraPlannerCfg), "neural": (NeuralPlanner, NeuralPlannerCfg), + "curobo": (CuroboPlanner, CuroboPlannerCfg), } def __init__(self, cfg: MotionGenCfg) -> None: diff --git a/examples/sim/motion_gen/curobo_motion_gen.py b/examples/sim/motion_gen/curobo_motion_gen.py new file mode 100644 index 00000000..c9dcf445 --- /dev/null +++ b/examples/sim/motion_gen/curobo_motion_gen.py @@ -0,0 +1,225 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""cuRobo V2 collision-aware motion-generation demo. + +Builds a single-arm Franka Panda with a static cuboid obstacle (mirrored in +both DexSim and the cuRobo collision world), plans a collision-free +end-effector move through the EmbodiChain ``MotionGenerator`` API, and replays +the returned full-DoF trajectory in the simulator. + +Requirements: an NVIDIA CUDA device and cuRobo V2 installed with matching +CUDA/PyTorch extras. See: +https://nvlabs.github.io/curobo/latest/getting-started/installation.html +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + +from embodichain import data as _data +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.objects import RigidObjectCfg +from embodichain.lab.sim.robots import FrankaPandaCfg +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenOptions, + MotionGenerator, + MoveType, + PlanState, +) +from embodichain.lab.sim.planners.curobo_planner import ( + CuroboPlanOptions, + CuroboPlannerCfg, + CuroboRobotProfileCfg, + CuroboWorldCfg, +) + +ROBOT_UID = "curobo_franka" +CONTROL_PART = "arm" +DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] +DEMO_BLOCK_POS = [0.45, 0.0, 0.18] +CUROBO_INSTALL_URL = ( + "https://nvlabs.github.io/curobo/latest/getting-started/installation.html" +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="cuRobo V2 motion-gen demo") + parser.add_argument( + "--headless", + action="store_true", + help="Run without opening the viewer window.", + ) + parser.add_argument( + "--step-repeat", + type=int, + default=4, + help="Simulation updates per planned waypoint during playback.", + ) + parser.add_argument( + "--hold-steps", + type=int, + default=20, + help="Simulation updates to hold before and after playback.", + ) + parser.add_argument( + "--no-warmup", + action="store_true", + help="Skip cuRobo planner warmup.", + ) + return parser.parse_args() + + +def _check_runtime() -> None: + """Fail fast with actionable guidance when CUDA/cuRobo are unavailable.""" + if not torch.cuda.is_available(): + raise RuntimeError( + "cuRobo V2 requires a CUDA-capable NVIDIA GPU, but none is " + "available. This demo cannot run on CPU." + ) + try: + import curobo # noqa: F401 + except ImportError as exc: + raise ImportError( + "cuRobo V2 is not installed. Install it with NVIDIA's CUDA-matched " + "extras, e.g. `pip install .[cu12]` or `pip install .[cu13]` " + f"(also `.[cu12-torch]` / `.[cu13-torch]`). See {CUROBO_INSTALL_URL}." + ) from exc + + +def _demo_world_path() -> str: + return str( + Path(_data.__file__).parent + / "assets" + / "curobo" + / "collision_franka_demo.yml" + ) + + +def _franka_profile() -> CuroboRobotProfileCfg: + """Explicit sim->cuRobo joint mapping; no index order is assumed.""" + sim_to_curobo = {f"fr3_joint{i}": f"panda_joint{i}" for i in range(1, 8)} + return CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names=sim_to_curobo, + fixed_joint_positions={ + "panda_finger_joint1": 0.04, + "panda_finger_joint2": 0.04, + }, + base_link_name="panda_link0", + tool_frame_name="panda_hand", + ) + + +def _build_scene(args: argparse.Namespace): + sim = SimulationManager( + SimulationManagerCfg( + headless=args.headless, + sim_device="cuda", + num_envs=1, + arena_space=2.0, + ) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) + ) + # Mirror the cuRobo cuboid in DexSim so planner and simulator agree. + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="demo_block", + shape=CubeCfg(size=DEMO_BLOCK_DIMS), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=DEMO_BLOCK_POS, + init_rot=[0.0, 0.0, 0.0], + ) + ) + return sim, robot + + +def _target_beyond_block(robot) -> torch.Tensor: + """An end-effector target that requires routing around the cuboid.""" + qpos = robot.get_qpos(name=CONTROL_PART) + fk = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True) + target = fk[0].clone() + target[:3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) + return target + + +def _play(sim, robot, trajectory: torch.Tensor, step_repeat: int) -> None: + all_ids = list(range(robot.dof)) + for w in range(trajectory.shape[1]): + robot.set_qpos(qpos=trajectory[:, w], joint_ids=all_ids) + sim.update(step=step_repeat) + + +def main() -> None: + args = parse_args() + _check_runtime() + + sim, robot = _build_scene(args) + if not args.headless: + sim.open_window() + sim.update(step=args.hold_steps) + + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg(world_config_path=_demo_world_path()), + warmup=not args.no_warmup, + ) + mg = MotionGenerator(MotionGenCfg(planner_cfg=cfg)) + + start_qpos = robot.get_qpos(name=CONTROL_PART) + target = _target_beyond_block(robot) + + result = mg.generate( + [PlanState.from_xpos(target.unsqueeze(0))], + MotionGenOptions( + start_qpos=start_qpos, + control_part=CONTROL_PART, + plan_opts=CuroboPlanOptions(control_part=CONTROL_PART), + ), + ) + + print(f"cuRobo success: {bool(result.success.item())}") + print(f"positions shape: {tuple(result.positions.shape)}") + print(f"duration: {float(result.duration[0]):.3f}s") + + if not bool(result.success.item()): + sim.destroy() + raise RuntimeError("cuRobo failed to find a collision-free trajectory.") + + # Replay the full-DoF trajectory. + _play(sim, robot, result.positions, step_repeat=max(args.step_repeat, 1)) + sim.update(step=args.hold_steps) + + final_q = result.positions[0:1, -1, :].to(robot.device) + fk = robot.compute_fk(qpos=final_q, name=CONTROL_PART, to_matrix=True) + err = float(torch.norm(fk[0, :3, 3] - target[:3, 3])) + print(f"final Cartesian position error: {err:.4f} m") + + sim.destroy() + + +if __name__ == "__main__": + main() diff --git a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py new file mode 100644 index 00000000..15165960 --- /dev/null +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -0,0 +1,169 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""Optional DexSim end-to-end test for cuRobo through AtomicActionEngine. + +Skipped when cuRobo or CUDA is unavailable. When both are present, it builds a +single-arm Franka + static cuboid scene, executes ``MoveEndEffector`` with +``planner_type='curobo'`` through the engine, and asserts a full-DoF +collision-aware trajectory that reaches the target after playback. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +# Module-level guards before any cuRobo-only import. +pytest.importorskip("curobo") +if not torch.cuda.is_available(): + pytest.skip("cuRobo V2 requires CUDA", allow_module_level=True) + +from embodichain import data as _data # noqa: E402 +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg # noqa: E402 +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 +from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 +from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 +from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 +from embodichain.lab.sim.planners import ( # noqa: E402 + MotionGenCfg, + MotionGenerator, +) +from embodichain.lab.sim.planners.curobo_planner import ( # noqa: E402 + CuroboPlannerCfg, + CuroboRobotProfileCfg, + CuroboWorldCfg, +) +from embodichain.lab.sim.atomic_actions import AtomicActionEngine # noqa: E402 +from embodichain.lab.sim.atomic_actions.actions import ( # noqa: E402 + MoveEndEffector, + MoveEndEffectorCfg, +) +from embodichain.lab.sim.atomic_actions.core import EndEffectorPoseTarget # noqa: E402 + +ROBOT_UID = "curobo_franka" +CONTROL_PART = "arm" +DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] +DEMO_BLOCK_POS = [0.45, 0.0, 0.18] +POS_TOL = 0.02 + + +def _demo_world_path() -> str: + return str( + Path(_data.__file__).parent + / "assets" + / "curobo" + / "collision_franka_demo.yml" + ) + + +def _franka_profile() -> CuroboRobotProfileCfg: + sim_to_curobo = {f"fr3_joint{i}": f"panda_joint{i}" for i in range(1, 8)} + return CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names=sim_to_curobo, + fixed_joint_positions={ + "panda_finger_joint1": 0.04, + "panda_finger_joint2": 0.04, + }, + base_link_name="panda_link0", + tool_frame_name="panda_hand", + ) + + +def _make_franka_curobo_engine(): + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cuda", num_envs=1) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) + ) + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="demo_block", + shape=CubeCfg(size=DEMO_BLOCK_DIMS), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=DEMO_BLOCK_POS, + init_rot=[0.0, 0.0, 0.0], + ) + ) + mg = MotionGenerator( + MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg(world_config_path=_demo_world_path()), + ) + ) + ) + engine = AtomicActionEngine(mg) + engine.register( + MoveEndEffector( + mg, + MoveEndEffectorCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part=CONTROL_PART, + sample_interval=80, + ), + ), + name="move_end_effector", + ) + return sim, robot, engine + + +def _reachable_target_beyond_demo_block(robot) -> torch.Tensor: + """A target beyond the cuboid so the planner must route around it.""" + qpos = robot.get_qpos(name=CONTROL_PART) + fk = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True) + target = fk[0].clone() + target[:3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) + return target + + +def _play_trajectory(sim, robot, trajectory: torch.Tensor, step_repeat: int = 1): + arm_ids = robot.get_joint_ids(name=CONTROL_PART) + for w in range(trajectory.shape[1]): + robot.set_qpos(qpos=trajectory[:, w], joint_ids=list(range(robot.dof))) + sim.update(step=step_repeat) + + +def _position_error(robot, target: torch.Tensor) -> float: + qpos = robot.get_qpos(name=CONTROL_PART) + fk = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True) + return float(torch.norm(fk[0, :3, 3] - target[:3, 3])) + + +@pytest.mark.requires_sim +@pytest.mark.slow +def test_atomic_move_end_effector_uses_curobo_v2(): + sim, robot, engine = _make_franka_curobo_engine() + try: + target = _reachable_target_beyond_demo_block(robot) + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + assert success.shape == (1,) + assert bool(success.item()) + assert trajectory.shape[2] == robot.dof + _play_trajectory(sim, robot, trajectory) + assert _position_error(robot, target) < POS_TOL + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() From 19000f8bd1c95c47071a647aeca6c7c016026592 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 11 Jul 2026 16:15:33 +0000 Subject: [PATCH 7/7] feat(planner): finalize curobo v2 integration Harden batched V2 collision worlds and add the documented atomic-action demo with tests. --- .../overview/sim/planners/curobo_planner.md | 200 ++++++ docs/source/overview/sim/planners/index.rst | 6 +- .../overview/sim/planners/motion_generator.md | 12 +- .../plans/2026-07-11-curobo-v2-integration.md | 19 +- ...2026-07-11-curobo-v2-integration-design.md | 22 +- .../assets/curobo/collision_franka_demo.yml | 2 +- .../lab/sim/atomic_actions/trajectory.py | 58 +- embodichain/lab/sim/planners/base_planner.py | 9 + .../lab/sim/planners/curobo_planner.py | 591 ++++++++++++++++-- .../lab/sim/planners/toppra_planner.py | 4 + examples/sim/motion_gen/curobo_motion_gen.py | 225 ------- examples/sim/planners/curobo_planner.py | 359 +++++++++++ tests/docs/test_curobo_planner_docs.py | 42 ++ tests/sim/atomic_actions/test_actions.py | 1 + .../test_curobo_motion_source_e2e.py | 33 +- .../test_trajectory_motion_source.py | 103 ++- tests/sim/planners/test_curobo_integration.py | 120 +++- tests/sim/planners/test_curobo_planner.py | 320 +++++++++- 18 files changed, 1765 insertions(+), 361 deletions(-) create mode 100644 docs/source/overview/sim/planners/curobo_planner.md delete mode 100644 examples/sim/motion_gen/curobo_motion_gen.py create mode 100644 examples/sim/planners/curobo_planner.py create mode 100644 tests/docs/test_curobo_planner_docs.py diff --git a/docs/source/overview/sim/planners/curobo_planner.md b/docs/source/overview/sim/planners/curobo_planner.md new file mode 100644 index 00000000..4b2c8769 --- /dev/null +++ b/docs/source/overview/sim/planners/curobo_planner.md @@ -0,0 +1,200 @@ +# cuRobo V2 Planner + +CuroboPlanner is EmbodiChain's optional, CUDA-accelerated and collision-aware +motion-planning backend. It implements the normal MotionGenerator and +atomic-action interfaces while cuRobo performs collision-aware inverse +kinematics and trajectory optimization. It supports Cartesian EEF_MOVE and +joint-space JOINT_MOVE requests for one configured control part at a time. + +planner_type="curobo" selects this backend. cuRobo V2 is deliberately not an +EmbodiChain core dependency: importing EmbodiChain planners does not import +cuRobo, and constructing this planner requires a CUDA-capable NVIDIA GPU. + +## Install cuRobo V2 + +Follow [NVIDIA's official cuRobo installation +guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html) +to install the V2 release that matches the CUDA driver and PyTorch environment. +The official flow clones cuRobo and uses a CUDA-matched extra: + +~~~bash +git clone https://github.com/NVlabs/curobo.git +cd curobo +uv venv --python 3.11 +source .venv/bin/activate + +# Choose exactly one command for the installed CUDA/PyTorch environment. +uv pip install .[cu12] # CUDA 12.x when PyTorch is already installed +uv pip install .[cu12-torch] # CUDA 12.x fresh environment, installs PyTorch +uv pip install .[cu13] # CUDA 13.x when PyTorch is already installed +uv pip install .[cu13-torch] # CUDA 13.x fresh environment, installs PyTorch + +python -c "import curobo; print(curobo.__version__)" +~~~ + +EmbodiChain does not install cuRobo transitively. Keep the cuRobo installation +in the same Python environment that runs the simulator, and use NVIDIA's +instructions when the CUDA or PyTorch version differs from this example. + +## Configure a control part + +Create one CuroboRobotProfileCfg for each EmbodiChain control part that may use +cuRobo. sim_to_curobo_joint_names is required: it maps the simulator's joint +names to the names in the cuRobo V2 robot profile, so no numeric joint ordering +is assumed. Lock non-controlled joints in the cuRobo robot profile itself so +they are not exposed in the loaded planner's active joint list. To plan a +gripper or another extra active joint, define a control part that includes it. +The retained simulator value of every such joint must equal the corresponding +cuRobo V2 `lock_joints` value throughout planning and playback. Atomic actions +preserve non-control joints in their full-DoF trajectory; they do not infer or +drive the profile's locked joints. For example, the stock Panda V2 profile +locks both fingers at `0.04`, so use the same simulated finger state or include +the fingers in the planned control part. A mismatch means cuRobo validates a +different collision geometry from the one replayed in DexSim. + +~~~python +from embodichain.lab.sim.planners import ( + CuroboPlannerCfg, + CuroboRobotProfileCfg, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, +) + +franka_profile = CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={ + f"fr3_joint{index}": f"panda_joint{index}" for index in range(1, 8) + }, + base_link_name="panda_link0", + tool_frame_name="panda_hand", + # DexSim's Panda TCP is 103.4 mm along +Z from cuRobo's panda_hand. + tool_frame_to_tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.1034], + [0.0, 0.0, 0.0, 1.0], + ], +) + +planner_cfg = CuroboPlannerCfg( + robot_uid="my_franka", + planner_type="curobo", + robot_profiles={"arm": franka_profile}, + world=CuroboWorldCfg(world_config_path="path/to/collision_world.yml"), +) +motion_generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg)) +~~~ + +The robot configuration must be a cuRobo V2 robot profile with collision +spheres and self-collision data. Generate or update that profile with V2 +RobotBuilder; a plain URDF alone is not sufficient for collision planning. The +mapped simulator joints must match the selected control part, and +tool_frame_name must name the cuRobo end-effector frame. + +`base_link_name`, when supplied, is checked against the loaded cuRobo model. +The adapter automatically rebases simulator-world Cartesian goals and dynamic +obstacle poses through the live simulator control-part base, so parallel arena +offsets and a moved robot base are handled. If the simulator and cuRobo base +frames use different fixed conventions, set `sim_base_to_curobo_base` to the +transform from the simulator base to the cuRobo base. Static collision YAML is +always authored in the cuRobo base/world frame. `tool_frame_to_tcp` is +different: it converts an EmbodiChain TCP goal into the chosen cuRobo tool +frame. Omit it only when both frames are identical. By convention, the adapter +uses `T_curobo,X = T_curobo,sim_base @ inv(T_world,sim_base) @ T_world,X`. +It obtains the simulator base from the control part's IK solver root; if that +part intentionally has no local solver, provide `sim_base_link_name` in the +profile instead. + +`CuroboPlannerCfg.use_cuda_graph` defaults to `False` for the same DexSim GPU +stream-safety reason. Enable it explicitly only after validating the local +simulation stack. + +CuroboWorldCfg.world_config_path names an explicit collision world. The initial +release accepts cuRobo cuboid, mesh, and voxel geometry. If obstacle poses +change at runtime, declare their names in +CuroboWorldCfg.dynamic_obstacle_names, provision +CuroboWorldCfg.collision_cache before planning, and pass their batched +(B, 4, 4) poses through CuroboPlanOptions.dynamic_obstacle_poses. Geometry is +not extracted automatically from DexSim. With the default shared world +(`multi_env=False`), all batch rows must provide the same obstacle pose; set +`multi_env=True` when each environment needs its own collision-world instance +(for example, different dynamic obstacle poses). In that mode a single mapping +YAML (including the supplied demo scene) is cloned into one V2 scene per batch +row. A top-level YAML list may instead define one mapping per row; it must have +either one entry (cloned) or exactly the active batch size. An empty configured +world is likewise materialized once per row so its per-environment cache is +allocated. Dynamic pose updates still require the named geometry to already +exist in every scene; the adapter does not insert new geometry at runtime. + +## Generate a motion + +MotionGenerator passes start_qpos and control_part to the cuRobo backend. For +Cartesian goals, leave EmbodiChain pre-interpolation disabled: cuRobo must +receive the original pose and preserves its own collision-checked samples. + +~~~python +import torch + +from embodichain.lab.sim.planners import MotionGenOptions, PlanState +from embodichain.lab.sim.planners.curobo_planner import CuroboPlanOptions + +goal_pose = torch.eye(4, device=robot.device).unsqueeze(0) +goal_pose[:, :3, 3] = torch.tensor( + [[0.55, 0.30, 0.45]], device=robot.device +) +result = motion_generator.generate( + [PlanState.from_xpos(goal_pose)], + MotionGenOptions( + start_qpos=robot.get_qpos(name="arm"), + control_part="arm", + plan_opts=CuroboPlanOptions(), + ), +) +assert result.success.all() +~~~ + +## Atomic actions and supported scope + +Single-arm MoveEndEffector is supported through the normal +motion_source="motion_gen" route. MoveJoints can opt in to collision-aware +joint-space planning with motion_source="motion_gen" and +planner_type="curobo". Movement phases of PickUp, Place, Press, and +MoveHeldObject can use the same single-arm static-world route. + +This first release intentionally has the following limits: + +- Only one configured control part is planned per request; coordinated dual-arm + planning and CoordinatedPickment are unsupported. +- Static cuboid, mesh, and voxel worlds plus named dynamic pose updates are + supported. Automatic scene extraction, arbitrary geometry insertion, and + removal are unsupported. +- A static collision YAML is for a fixed-base robot. With a moving base, + publish each relevant world obstacle as a named dynamic pose for every plan; + automatic reprojection of static YAML obstacles is unsupported. +- attached-object collision geometry, automatic attachment/detachment, and + collision-aware carrying of a held object are unsupported. +- Non-control joints must remain at the matching cuRobo V2 `lock_joints` + values. The adapter does not yet validate cross-model locked-joint name/value + equivalence automatically. +- The legacy Gym ActionBank path is unsupported. +- CPU execution and cuRobo V1 compatibility are unsupported. + +## Demo + +After installing cuRobo V2 and configuring a CUDA simulation environment, run +the Panda obstacle-avoidance demo from the repository root: + +~~~bash +python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1 +~~~ + +The demo mirrors a cuboid in DexSim and the cuRobo collision world, prints the +result status and trajectory shape, then replays the returned full-DoF +trajectory. It disables cuRobo CUDA graph capture by default because graph +capture can conflict with DexSim GPU physics; pass `--cuda-graph` only after +validating that the local simulator stream setup supports it. Headless runs +automatically record this fixed offscreen camera view to an MP4. Set an explicit +destination with `--record-save-path outputs/videos/curobo_demo.mp4`, adjust +the rate with `--record-fps`, or pass `--disable-record` to skip recording. See +[MotionGenerator](motion_generator.md) for the common planner interface. diff --git a/docs/source/overview/sim/planners/index.rst b/docs/source/overview/sim/planners/index.rst index d1320be1..9453939f 100644 --- a/docs/source/overview/sim/planners/index.rst +++ b/docs/source/overview/sim/planners/index.rst @@ -19,12 +19,13 @@ Overview The `embodichain` project provides a unified interface for robot trajectory planning, supporting both joint space and Cartesian space interpolation. The main planners include: -- **MotionGenerator**: A unified trajectory planning interface that supports joint/Cartesian interpolation, automatic constraint handling, flexible planner selection, and is easily extensible for collision checking and additional planners. +- **MotionGenerator**: A unified trajectory planning interface that supports joint/Cartesian interpolation, automatic constraint handling, flexible planner selection, and backend-specific collision-aware planning. - **ToppraPlanner**: A time-optimal trajectory planner based on the TOPPRA library, supporting joint trajectory generation under velocity and acceleration constraints. - **NeuralPlanner** (experimental): A learning-based EEF waypoint planner for Franka Panda. +- **CuroboPlanner** (optional): A CUDA-only cuRobo V2 backend for collision-aware single-arm Cartesian and joint-space planning. - **TrajectorySampleMethod**: An enumeration for trajectory sampling strategies, supporting sampling by time, quantity, or distance. -These tools can be used to generate smooth and dynamically feasible robot trajectories, and are extensible for future collision checking and various sampling requirements. +These tools can be used to generate smooth and dynamically feasible robot trajectories. Install cuRobo separately when collision-aware planning against an explicit cuRobo world is required. Use NeuralPlanner (experimental) when you have a trained APG checkpoint and need learned EEF waypoint rollout on Franka Panda. @@ -37,5 +38,6 @@ See also toppra_planner.md neural_planner.md + curobo_planner.md trajectory_sample_method.md motion_generator.md diff --git a/docs/source/overview/sim/planners/motion_generator.md b/docs/source/overview/sim/planners/motion_generator.md index 4d8b53db..6a875ced 100644 --- a/docs/source/overview/sim/planners/motion_generator.md +++ b/docs/source/overview/sim/planners/motion_generator.md @@ -1,13 +1,13 @@ # MotionGenerator -`MotionGenerator` provides a unified interface for robot trajectory planning, supporting both joint space and Cartesian space interpolation. It is designed to work with different planners (such as ToppraPlanner) and can be extended to support collision checking in the future. +`MotionGenerator` provides a unified interface for robot trajectory planning, supporting both joint space and Cartesian space interpolation. It dispatches to the selected backend: TOPPRA and NeuralPlanner retain their existing behavior, while the optional cuRobo V2 backend performs collision-aware planning against an explicit cuRobo world. ## Features -* **Unified planning interface**: Supports trajectory planning with or without collision checking (collision checking is reserved for future implementation). -* **Flexible planner selection**: Supports TOPPRA and NeuralPlanner (experimental). +* **Unified planning interface**: Supports interpolation-oriented planners and collision-aware cuRobo V2 planning through one `generate()` API. +* **Flexible planner selection**: Supports TOPPRA, NeuralPlanner (experimental), and the optional CUDA-only CuroboPlanner backend. * **Automatic constraint handling**: Retrieves velocity and acceleration limits from the robot or uses user-specified/default values. -* **Supports both joint and Cartesian interpolation**: Generates discrete trajectories using either joint space or Cartesian space interpolation. +* **Backend-aware target handling**: Generates discrete trajectories using joint or Cartesian interpolation where appropriate; cuRobo receives original Cartesian goals so it can perform collision-aware IK itself. * **Convenient sampling**: Supports various sampling strategies via `TrajectorySampleMethod`. ## Usage @@ -181,5 +181,7 @@ print(f"Estimated sample count: {sample_count}") * The planner type can be specified as a string or `PlannerType` enum. * If the robot provides its own joint limits, those will be used; otherwise, default or user-specified limits are applied. * For Cartesian interpolation, inverse kinematics (IK) is used to compute joint configurations for each interpolated pose. -* The class is designed to be extensible for additional planners and collision checking in the future. +* Backends declare whether pre-interpolation is safe and whether their returned samples must be preserved. cuRobo V2 disables EmbodiChain Cartesian pre-interpolation and preserves its collision-checked samples. +* CuroboPlanner is optional and requires CUDA plus a matching cuRobo V2 installation; see [the cuRobo planner page](curobo_planner.md) and [NVIDIA's installation guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html). +* Run the collision-aware Panda demo with `python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1`. * The sample count estimation is useful for predicting computational load and memory requirements. diff --git a/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md b/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md index 27480f21..0cf38945 100644 --- a/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md +++ b/docs/superpowers/plans/2026-07-11-curobo-v2-integration.md @@ -285,15 +285,17 @@ class CuroboRobotProfileCfg: robot_config_path: str = MISSING sim_to_curobo_joint_names: dict[str, str] = MISSING active_joint_names: list[str] | None = None - fixed_joint_positions: dict[str, float] = {} base_link_name: str | None = None + sim_base_link_name: str | None = None + sim_base_to_curobo_base: list[list[float]] | None = None tool_frame_name: str | None = None + tool_frame_to_tcp: list[list[float]] | None = None @configclass class CuroboWorldCfg: world_config_path: str | None = None - collision_cache: dict[str, int] = {"cuboid": 8, "mesh": 2, "voxel": 1} + collision_cache: dict[str, int] = {"cuboid": 8, "mesh": 2} dynamic_obstacle_names: list[str] = [] multi_env: bool = False @@ -307,7 +309,7 @@ class CuroboPlannerCfg(BasePlannerCfg): collision_activation_distance: float = 0.01 max_attempts: int = 5 max_planning_time: float | None = None - use_cuda_graph: bool = True + use_cuda_graph: bool = False interpolation_dt: float = 0.025 @@ -428,7 +430,7 @@ class CuroboPlanner(BasePlanner): return self._plan_segments(target_states, start, control_part, backend, options) ``` -`_get_backend` must reject a non-CUDA robot device; invoke V2 `MotionPlannerCfg.create` with the profile path, `scene_model=world_config_path`, non-`None` `collision_cache`, `max_batch_size=batch_size`, `multi_env=cfg.world.multi_env`, `optimizer_collision_activation_distance`, `use_cuda_graph`, and `interpolation_dt`. Instantiate `MotionPlanner` for `B == 1` and `BatchMotionPlanner` for `B > 1`; call V2 warmup once per cache entry when `cfg.warmup` is true. Cache only entries matching the actual batch size. When `profile.active_joint_names` is set, compare it with `backend.planner.joint_names` and raise on any missing, duplicate, or differently ordered name. +`_get_backend` must reject a non-CUDA robot device; invoke V2 `MotionPlannerCfg.create` with the profile path, non-`None` `collision_cache`, `max_batch_size=batch_size`, `multi_env=cfg.world.multi_env`, `optimizer_collision_activation_distance`, `use_cuda_graph`, and `interpolation_dt`. In multi-environment mode it must materialize exactly `batch_size` scene mappings: clone a singleton mapping/empty world, or accept a top-level YAML list with exactly that many mappings. V2 infers collision-world count from this list, not from `max_batch_size`. Instantiate `MotionPlanner` for `B == 1` and `BatchMotionPlanner` for `B > 1`; call V2 warmup once per cache entry when `cfg.warmup` is true. Cache only entries matching the actual batch size. When `profile.active_joint_names` is set, compare it with `backend.planner.joint_names` and raise on any missing, duplicate, or differently ordered name. For each segment, construct V2 states and goals as follows: @@ -629,7 +631,7 @@ Add `embodichain/data/assets/curobo/collision_franka_demo.yml`: ```yaml cuboid: - - name: demo_block + demo_block: dims: [0.18, 0.40, 0.36] pose: [0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0] ``` @@ -689,9 +691,14 @@ FRANKA_SIM_TO_CUROBO = { profile = CuroboRobotProfileCfg( robot_config_path="franka.yml", sim_to_curobo_joint_names=FRANKA_SIM_TO_CUROBO, - fixed_joint_positions={"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}, base_link_name="panda_link0", tool_frame_name="panda_hand", + tool_frame_to_tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.1034], + [0.0, 0.0, 0.0, 1.0], + ], ) ``` diff --git a/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md b/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md index 17712a18..6e90d6f7 100644 --- a/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md +++ b/docs/superpowers/specs/2026-07-11-curobo-v2-integration-design.md @@ -46,7 +46,7 @@ module import time. - Cartesian pose planning and joint-space planning through the existing `PlanState` / `PlanResult` interface. - Single-arm control-part planning, including robust named-joint ordering, - fixed non-controlled joints, TCP, root-frame, and joint-limit handling. + locked non-controlled joints, TCP, root-frame, and joint-limit handling. - Batch-aware planning: the adapter accepts EmbodiChain's leading batch dimension and uses cuRobo's V2 batch planner where multiple environments are requested. The CUDA integration test establishes the single-environment @@ -83,9 +83,11 @@ The planner package adds focused, serializable config objects: CuroboRobotProfileCfg robot_config_path active_joint_names - fixed_joint_positions base_link_name + sim_base_link_name + sim_base_to_curobo_base tool_frame_name + tool_frame_to_tcp CuroboWorldCfg world_config_path @@ -113,8 +115,14 @@ CuroboPlanOptions(PlanOptions) Each `robot_profiles` key is an EmbodiChain control-part name. The selected profile explicitly maps between simulator joint names and cuRobo joint names; the backend never assumes that simulator indices and cuRobo indices are equal. -The profile also pins non-controlled joints, including gripper joints, to -current or configured values when the cuRobo robot model contains them. +Any non-controlled joints, including gripper joints, must be locked in the +cuRobo V2 robot profile so they do not appear in the backend's active joint +list. Their preserved simulator values must equal the V2 profile's +`lock_joints` values during planning and atomic-action playback; the first +release documents this cross-model lock contract but does not automatically +validate joint-name/value equivalence. The backend rejects an active-joint/ +profile mismatch rather than planning collision geometry that EmbodiChain will +not execute. `CuroboPlanner` is a `BasePlanner` implementation. It lazily imports V2 types, creates and warms a `MotionPlanner` for a one-environment request and a @@ -142,6 +150,8 @@ move types fail before planning with a clear `ValueError`. `BasePlanner` gains two explicit extension points: - `preinterpolate_targets: bool`, which defaults to `True`; +- `supports_joint_move: bool`, which defaults to `False` and keeps + Cartesian-only backends on local joint interpolation for atomic phases; - `with_motion_context(options, *, start_qpos, control_part)`, whose base implementation returns the supplied options unchanged. @@ -172,7 +182,9 @@ than a fragile implicit scene scan. All obstacle poses, goal poses, robot base poses, and tool poses use a single documented world coordinate convention. The adapter converts EmbodiChain 4x4 matrices to cuRobo's position/quaternion representation at this boundary -and tests the quaternion ordering directly. +and tests the quaternion ordering directly. Static collision YAML is authored +in the cuRobo base frame and therefore applies to fixed-base scenes; a moving +base must publish relevant obstacles through named dynamic updates. ### Atomic-Action Integration diff --git a/embodichain/data/assets/curobo/collision_franka_demo.yml b/embodichain/data/assets/curobo/collision_franka_demo.yml index 3503d4b4..1ced97d3 100644 --- a/embodichain/data/assets/curobo/collision_franka_demo.yml +++ b/embodichain/data/assets/curobo/collision_franka_demo.yml @@ -6,6 +6,6 @@ # # Pose convention: [x, y, z, qw, qx, qy, qz]. cuboid: - - name: demo_block + demo_block: dims: [0.18, 0.40, 0.36] pose: [0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0] diff --git a/embodichain/lab/sim/atomic_actions/trajectory.py b/embodichain/lab/sim/atomic_actions/trajectory.py index a6b217b0..ddd9fc71 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory.py +++ b/embodichain/lab/sim/atomic_actions/trajectory.py @@ -34,6 +34,14 @@ from embodichain.lab.sim.planners import MotionGenerator +def _resolve_runtime_device(device: torch.device | str) -> torch.device: + """Resolve an indexless CUDA device to the active concrete GPU index.""" + resolved = torch.device(device) + if resolved.type == "cuda" and resolved.index is None: + return torch.device(f"cuda:{torch.cuda.current_device()}") + return resolved + + class TrajectoryBuilder: """Stateless trajectory utilities shared by every atomic action. @@ -45,7 +53,7 @@ class TrajectoryBuilder: def __init__(self, motion_generator: MotionGenerator) -> None: self.motion_generator = motion_generator self.robot = motion_generator.robot - self.device = self.robot.device + self.device = _resolve_runtime_device(self.robot.device) # ------------------------------------------------------------------ # Success / shape helpers @@ -531,10 +539,11 @@ def plan_joint_motion( ) -> tuple[torch.Tensor, torch.Tensor]: """Plan a joint-space trajectory through one or more target waypoints. - For ``motion_source='motion_gen'`` this delegates to the same - MotionGenerator validation path as Cartesian planning, building batched - ``JOINT_MOVE`` PlanStates. For ``motion_source='ik_interp'`` (the - default) it returns an all-success linear interpolation. + For ``motion_source='motion_gen'``, this delegates only when the + selected backend advertises ``supports_joint_move``. Cartesian-only + backends (such as the neural planner) retain the deterministic local + interpolation for joint-only phases. ``motion_source='ik_interp'`` + always uses that local interpolation. Returns: ``(success:(B,), trajectory:(B, N, arm_dof))``. @@ -549,25 +558,26 @@ def plan_joint_motion( ValueError, ) self._validate_planner_type(cfg) - if target_qpos.dim() == 2: - target_qpos = target_qpos.unsqueeze(1) # (B, 1, D) - plan_states = [ - PlanState(qpos=target_qpos[:, j], move_type=MoveType.JOINT_MOVE) - for j in range(target_qpos.shape[1]) - ] - plan_opts = self._build_plan_opts(cfg, n_waypoints) - result: PlanResult = self.motion_generator.generate( - plan_states, - options=MotionGenOptions( - start_qpos=start_qpos, - control_part=control_part, - plan_opts=plan_opts, - is_interpolate=self.motion_generator.planner.preinterpolate_targets, - ), - ) - return self._process_motion_gen_result( - result, start_qpos, n_waypoints, arm_dof - ) + if self.motion_generator.planner.supports_joint_move: + if target_qpos.dim() == 2: + target_qpos = target_qpos.unsqueeze(1) # (B, 1, D) + plan_states = [ + PlanState(qpos=target_qpos[:, j], move_type=MoveType.JOINT_MOVE) + for j in range(target_qpos.shape[1]) + ] + plan_opts = self._build_plan_opts(cfg, n_waypoints) + result: PlanResult = self.motion_generator.generate( + plan_states, + options=MotionGenOptions( + start_qpos=start_qpos, + control_part=control_part, + plan_opts=plan_opts, + is_interpolate=self.motion_generator.planner.preinterpolate_targets, + ), + ) + return self._process_motion_gen_result( + result, start_qpos, n_waypoints, arm_dof + ) success = torch.ones(start_qpos.shape[0], dtype=torch.bool, device=self.device) trajectory = self.plan_joint_traj(start_qpos, target_qpos, n_waypoints) return success, trajectory diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index cda8f011..45612bd3 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -172,6 +172,15 @@ def __init__(self, cfg: BasePlannerCfg): waypoint count. """ + supports_joint_move: bool = False + """Whether the backend accepts :attr:`MoveType.JOINT_MOVE` targets. + + Atomic actions use this capability to decide whether their joint-only + phases may be delegated through ``MotionGenerator``. Cartesian-only + planners retain the deterministic local joint interpolation for those + phases. + """ + def default_plan_options(self) -> PlanOptions: """Return backend-default planning options.""" return PlanOptions() diff --git a/embodichain/lab/sim/planners/curobo_planner.py b/embodichain/lab/sim/planners/curobo_planner.py index bb404e65..8cdb29b5 100644 --- a/embodichain/lab/sim/planners/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo_planner.py @@ -29,11 +29,14 @@ from __future__ import annotations import importlib +from copy import deepcopy from dataclasses import MISSING, dataclass +from pathlib import Path from types import SimpleNamespace from typing import TYPE_CHECKING import torch +import yaml from embodichain.utils import configclass, logger from embodichain.utils.math import quat_from_matrix @@ -71,8 +74,13 @@ class CuroboRobotProfileCfg: Each ``CuroboPlannerCfg.robot_profiles`` key is an EmbodiChain control-part name. The profile explicitly maps simulator joint names to cuRobo joint names so the backend never assumes simulator and cuRobo joint indices agree. - Non-controlled joints (e.g. gripper joints) present in the cuRobo model are - pinned to ``fixed_joint_positions``. + The initial release requires the loaded cuRobo planner's active joints to + match the selected control part exactly. Lock non-controlled joints (for + example gripper joints) in the cuRobo V2 robot profile so they are not + exposed as active planner joints. The simulator values of those joints + must remain equal to the V2 profile's ``lock_joints`` values while a plan + is executed; the adapter intentionally preserves non-control simulator + joints in the full-DoF atomic-action output. """ robot_config_path: str = MISSING @@ -84,15 +92,42 @@ class CuroboRobotProfileCfg: active_joint_names: list[str] | None = None """Optional explicit cuRobo active-joint ordering, validated against the backend.""" - fixed_joint_positions: dict[str, float] = {} - """cuRobo joint names pinned to a constant value (e.g. gripper finger joints).""" - base_link_name: str | None = None - """cuRobo robot base link name.""" + """Expected cuRobo robot base link, validated against the loaded V2 model. + + This is a consistency check. Use ``sim_base_to_curobo_base`` when the + simulator control-part base and this V2 base use different fixed frames. + Static collision YAML remains authored in the cuRobo base/world frame. + """ + + sim_base_to_curobo_base: list[list[float]] | None = None + """Fixed transform from the simulator control-part base to the cuRobo base. + + The adapter uses this transform together with the live simulator base pose + to convert simulator-world Cartesian goals and dynamic obstacle poses into + cuRobo's base frame. ``None`` means the two base frames coincide. + """ + + sim_base_link_name: str | None = None + """Simulator link physically equivalent to the control-part base. + + When omitted, the adapter uses the EmbodiChain solver's ``root_link_name``. + Set this explicitly for a cuRobo-planned control part that has no local + EmbodiChain IK solver. + """ tool_frame_name: str | None = None """cuRobo tool frame name used as the planning target.""" + tool_frame_to_tcp: list[list[float]] | None = None + """Fixed transform from the cuRobo tool frame to the simulator TCP frame. + + EmbodiChain Cartesian targets are expressed in the simulator TCP frame. + When that frame is not identical to ``tool_frame_name``, provide this + homogeneous transform so the adapter can convert the target before it is + sent to cuRobo. ``None`` means the two frames are identical. + """ + @configclass class CuroboWorldCfg: @@ -101,14 +136,29 @@ class CuroboWorldCfg: world_config_path: str | None = None """Path/identifier of a cuRobo V2 scene profile (cuboid/mesh/voxel obstacles).""" - collision_cache: dict[str, int] = {"cuboid": 8, "mesh": 2, "voxel": 1} - """Per-geometry cache capacity created before world updates.""" + collision_cache: dict[str, int | dict[str, int | float | list[float]]] = { + "cuboid": 8, + "mesh": 2, + } + """Per-geometry cache capacity created before world updates. + + cuRobo V2 accepts integer ``cuboid`` and ``mesh`` capacities. A ``voxel`` + cache, when needed for dynamic voxel worlds, must instead be a dictionary + with V2's ``layers``, ``dims``, and ``voxel_size`` fields. + """ dynamic_obstacle_names: list[str] = [] """Obstacle names whose poses may be updated between plans.""" multi_env: bool = False - """Whether the cuRobo world is shared across multiple environments.""" + """Whether cuRobo allocates one collision-world instance per environment. + + ``False`` shares one world and therefore requires equal rebased dynamic + obstacle poses across batch rows. ``True`` materializes one V2 scene model + per batch row, supporting independently updated obstacle poses for each + environment. A mapping YAML is cloned for every row; a top-level YAML list + may instead provide one explicit mapping per row. + """ @configclass @@ -135,8 +185,12 @@ class CuroboPlannerCfg(BasePlannerCfg): max_planning_time: float | None = None """Post-plan validation budget (seconds). ``None`` skips the timing check.""" - use_cuda_graph: bool = True - """Whether cuRobo may use CUDA graphs internally.""" + use_cuda_graph: bool = False + """Whether cuRobo may use CUDA graphs internally. + + Disabled by default because CUDA graph capture can conflict with DexSim's + GPU physics stream. Enable only after validating the local stream setup. + """ interpolation_dt: float = 0.025 """Interpolation step (seconds) used by cuRobo and as a dt fallback.""" @@ -188,7 +242,11 @@ def _reorder_by_names( Raises: ValueError: If the two name sets are not equal as sets. """ - if sorted(from_names) != sorted(to_names): + if ( + len(from_names) != len(set(from_names)) + or len(to_names) != len(set(to_names)) + or sorted(from_names) != sorted(to_names) + ): raise ValueError( f"Cannot reorder joints: source names {from_names} and target names " f"{to_names} are not the same set." @@ -219,8 +277,12 @@ def _matrix_to_position_quaternion( f"Expected (B, 4, 4) pose matrices, got shape {tuple(matrix.shape)}." ) matrix = matrix.to(dtype=torch.float32) - position = matrix[:, :3, 3] - quaternion = quat_from_matrix(matrix[:, :3, :3]) # wxyz + # V2's Pose inverse/update kernels require contiguous float32 tensors. + # Column/rotation slices of a homogeneous transform are views with strides, + # so materialize them at the adapter boundary rather than relying on a + # caller-specific layout. + position = matrix[:, :3, 3].contiguous() + quaternion = quat_from_matrix(matrix[:, :3, :3]).contiguous() # wxyz return position, quaternion @@ -291,6 +353,7 @@ def _require_curobo() -> "Any": JointState=types_mod.JointState, Pose=types_mod.Pose, GoalToolPose=types_mod.GoalToolPose, + DeviceCfg=types_mod.DeviceCfg, ) @@ -324,6 +387,7 @@ class CuroboPlanner(BasePlanner): preinterpolate_targets = False preserve_plan_samples = True + supports_joint_move = True def __init__(self, cfg: CuroboPlannerCfg) -> None: super().__init__(cfg) @@ -431,6 +495,84 @@ def _resolve_start_qpos( # Backend construction / caching # ------------------------------------------------------------------ + def _materialize_multi_env_scene_model( + self, world_config_path: str | None, batch_size: int + ) -> list[dict]: + """Return exactly one cuRobo V2 scene mapping for every batch row. + + cuRobo V2 infers the collision-world count from the length of a scene + model list; ``multi_env=True`` and ``max_batch_size`` alone do not + allocate per-environment collision data. A single mapping YAML is + therefore cloned for every row. A top-level YAML list supports + explicitly different static worlds, but it must contain either one + mapping (cloned) or exactly ``batch_size`` mappings. + + Args: + world_config_path: cuRobo scene identifier/path, or ``None`` for + an initially empty collision world. + batch_size: Number of simultaneous planning environments. + + Returns: + A list of independent V2 scene mappings with length ``batch_size``. + + Raises: + ValueError: If the YAML cannot be loaded or has an incompatible + top-level structure/count. + """ + if batch_size < 1: + logger.log_error( + f"multi-env cuRobo batch_size must be positive, got {batch_size}.", + ValueError, + ) + + if world_config_path is None: + # Even an initially empty collision world needs one scene mapping + # per row. Otherwise V2's SceneCollisionCfg defaults to num_envs=1 + # despite multi_env=True and later dynamic updates fail for row > 0. + return [{} for _ in range(batch_size)] + + scene_path = Path(world_config_path) + if not scene_path.is_absolute(): + content_mod = importlib.import_module("curobo.content") + scene_path = Path(content_mod.get_scene_configs_path()) / scene_path + try: + with scene_path.open(encoding="utf-8") as scene_file: + scene_model = yaml.safe_load(scene_file) + except (OSError, yaml.YAMLError) as exc: + logger.log_error( + f"Unable to load cuRobo V2 scene configuration " + f"'{world_config_path}': {exc}", + ValueError, + ) + raise AssertionError("unreachable") from exc + + if isinstance(scene_model, dict): + return [deepcopy(scene_model) for _ in range(batch_size)] + if isinstance(scene_model, list): + if not scene_model or not all( + isinstance(scene, dict) for scene in scene_model + ): + logger.log_error( + "A multi-env cuRobo scene YAML list must contain one or more " + "mapping worlds.", + ValueError, + ) + if len(scene_model) == 1: + return [deepcopy(scene_model[0]) for _ in range(batch_size)] + if len(scene_model) == batch_size: + return [deepcopy(scene) for scene in scene_model] + logger.log_error( + "A multi-env cuRobo scene YAML list must have one world to clone " + f"or exactly batch_size={batch_size} worlds; got {len(scene_model)}.", + ValueError, + ) + logger.log_error( + "A cuRobo V2 scene YAML must contain a mapping world or a list of " + f"mapping worlds, got {type(scene_model).__name__}.", + ValueError, + ) + raise AssertionError("unreachable") + def _get_backend( self, profile: CuroboRobotProfileCfg, @@ -443,14 +585,27 @@ def _get_backend( if key in self._backend_cache: return self._backend_cache[key] + sim_joint_names = self._resolve_sim_joint_names(control_part) world_cfg = self.cfg.world collision_cache = ( dict(world_cfg.collision_cache) if world_cfg.collision_cache else None ) + curobo_device = self.device + if curobo_device.index is None: + # Warp's V2 collision cache indexes CUDA devices by integer and + # cannot consume the indexless torch.device("cuda") used by + # DexSim's default CUDA configuration. + curobo_device = torch.device(f"cuda:{torch.cuda.current_device()}") + scene_model = world_cfg.world_config_path + if multi_env: + scene_model = self._materialize_multi_env_scene_model( + world_cfg.world_config_path, int(batch_size) + ) planner_cfg = self._bindings.MotionPlannerCfg.create( robot=profile.robot_config_path, - scene_model=world_cfg.world_config_path, + scene_model=scene_model, collision_cache=collision_cache, + device_cfg=self._bindings.DeviceCfg(device=curobo_device), max_batch_size=int(batch_size), multi_env=bool(multi_env), optimizer_collision_activation_distance=self.cfg.collision_activation_distance, @@ -462,27 +617,127 @@ def _get_backend( else: planner = self._bindings.BatchMotionPlanner(planner_cfg) - if profile.active_joint_names is not None: - expected = list(profile.active_joint_names) - actual = list(planner.joint_names) - if expected != actual: - logger.log_error( - f"active_joint_names {expected} do not match cuRobo model " - f"joints {actual} (missing/duplicate/out-of-order).", - ValueError, - ) + self._validate_profile_joint_names( + profile, sim_joint_names, list(planner.joint_names) + ) + self._validate_base_link_name(profile, planner) + tool_frame = self._resolve_tool_frame(profile, planner) backend = _CuroboBackend( planner=planner, - tool_frame=profile.tool_frame_name, + control_part=control_part, + sim_joint_names=sim_joint_names, + tool_frame=tool_frame, profile=profile, batch_size=int(batch_size), ) if self.cfg.warmup: - planner.warmup() + planner.warmup(enable_graph=bool(self.cfg.use_cuda_graph)) self._backend_cache[key] = backend return backend + def _resolve_sim_joint_names(self, control_part: str) -> list[str]: + """Return simulator control-part joints in the robot's canonical order.""" + control_parts = getattr(self.robot, "control_parts", None) + if not control_parts or control_part not in control_parts: + logger.log_error( + f"Robot '{self.cfg.robot_uid}' has no control part '{control_part}'. " + "cuRobo requires an explicit ordered control-part joint list.", + ValueError, + ) + return list(control_parts[control_part]) + + def _validate_profile_joint_names( + self, + profile: CuroboRobotProfileCfg, + sim_joint_names: list[str], + curobo_joint_names: list[str], + ) -> None: + """Validate the profile mapping before a CUDA planning call.""" + sim_to_curobo = profile.sim_to_curobo_joint_names + if set(sim_to_curobo) != set(sim_joint_names): + logger.log_error( + "sim_to_curobo_joint_names keys must exactly match the robot " + f"control-part joints {sim_joint_names}; got {list(sim_to_curobo)}.", + ValueError, + ) + mapped_names = [sim_to_curobo[name] for name in sim_joint_names] + if len(mapped_names) != len(set(mapped_names)): + logger.log_error( + "sim_to_curobo_joint_names maps multiple simulator joints to " + f"the same cuRobo joint: {mapped_names}.", + ValueError, + ) + missing = [name for name in mapped_names if name not in curobo_joint_names] + if missing: + logger.log_error( + "cuRobo profile is missing mapped active joints " + f"{missing}; planner joints are {curobo_joint_names}.", + ValueError, + ) + unmapped = [ + name for name in curobo_joint_names if name not in set(mapped_names) + ] + if unmapped: + logger.log_error( + "cuRobo planner exposes joints outside the requested control " + f"part: {unmapped}. Lock non-controlled joints in the V2 robot " + "profile or select a control part that includes them.", + ValueError, + ) + if profile.active_joint_names is not None: + expected = list(profile.active_joint_names) + if expected != curobo_joint_names: + logger.log_error( + f"active_joint_names {expected} do not match cuRobo model " + f"joints {curobo_joint_names} (missing/duplicate/out-of-order).", + ValueError, + ) + + def _resolve_tool_frame( + self, profile: CuroboRobotProfileCfg, planner: "Any" + ) -> str: + """Resolve and validate the single V2 tool frame used for pose goals.""" + tool_frames = list(getattr(planner, "tool_frames", [])) + tool_frame = profile.tool_frame_name + if tool_frame is None: + if len(tool_frames) != 1: + logger.log_error( + "tool_frame_name is required when the cuRobo profile exposes " + f"multiple tool frames: {tool_frames}.", + ValueError, + ) + return tool_frames[0] + if tool_frames and tool_frame not in tool_frames: + logger.log_error( + f"tool_frame_name '{tool_frame}' is not available in the cuRobo " + f"profile tool frames {tool_frames}.", + ValueError, + ) + return tool_frame + + def _validate_base_link_name( + self, profile: CuroboRobotProfileCfg, planner: "Any" + ) -> None: + """Ensure an explicitly configured base link matches the V2 model.""" + expected = profile.base_link_name + if expected is None: + return + kinematics = getattr(planner, "kinematics", None) + actual = getattr(kinematics, "base_link", None) + if actual is None: + logger.log_error( + "cuRobo planner did not expose a kinematics.base_link, so " + f"base_link_name={expected!r} cannot be validated.", + ValueError, + ) + if actual != expected: + logger.log_error( + f"CuroboRobotProfileCfg.base_link_name={expected!r} does not " + f"match the loaded cuRobo V2 base link {actual!r}.", + ValueError, + ) + # ------------------------------------------------------------------ # Segment planning # ------------------------------------------------------------------ @@ -508,6 +763,7 @@ def _plan_segments( current = start.clone() for seg_idx, target in enumerate(target_states): + self._validate_segment_batch(target, B, seg_idx) current_state = self._to_curobo_joint_state(current, backend) if target.move_type == MoveType.EEF_MOVE: if target.xpos is None: @@ -535,11 +791,19 @@ def _plan_segments( ValueError, ) - seg_success, seg_positions, seg_dt = self._extract_segment( - v2_result, backend - ) + if v2_result is None: + # V2 returns None when no seed reaches a valid solution. Keep + # the standard EmbodiChain failure contract instead of + # dereferencing a result that does not exist. + seg_success = torch.zeros(B, dtype=torch.bool, device=self.device) + seg_positions = current.unsqueeze(1) + seg_dt = torch.zeros(B, 1, dtype=torch.float32, device=self.device) + else: + seg_success, seg_positions, seg_dt = self._extract_segment( + v2_result, backend + ) seg_success = seg_success.to(self.device) & alive - if self.cfg.max_planning_time is not None: + if v2_result is not None and self.cfg.max_planning_time is not None: total_time = self._extract_total_time(v2_result, B) over = total_time > float(self.cfg.max_planning_time) seg_success = seg_success & (~over) @@ -562,6 +826,33 @@ def _plan_segments( return self._assemble_result(per_env_samples, per_env_dt, start, alive, B, D) + def _validate_segment_batch( + self, target: PlanState, start_batch_size: int, segment_index: int + ) -> None: + """Reject target batches that cannot pair with the planning start state.""" + if target.move_type == MoveType.EEF_MOVE: + values = target.xpos + expected_dims = (3,) + elif target.move_type == MoveType.JOINT_MOVE: + values = target.qpos + expected_dims = (1, 2) + else: + return + if values is None: + return + values = torch.as_tensor(values) + if values.dim() not in expected_dims: + # The type-specific conversion path will report the more useful + # shape error below; only check valid target shapes here. + return + target_batch_size = 1 if values.dim() == 1 else values.shape[0] + if target_batch_size != start_batch_size: + logger.log_error( + f"Segment {segment_index} target batch {target_batch_size} does " + f"not match planning start batch {start_batch_size}.", + ValueError, + ) + def _extract_segment( self, v2_result: "Any", backend: "_CuroboBackend" ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -596,20 +887,20 @@ def _extract_segment( if length < max_len: full[b, length:] = position[b, length - 1].float().to(self.device) - seg_positions = self._map_curobo_to_sim(full, traj.joint_names, backend.profile) - seg_dt = self._extract_dt(v2_result, traj, max_len, B) + seg_positions = self._map_curobo_to_sim(full, traj.joint_names, backend) + seg_dt = self._extract_dt(traj, last_tstep, max_len, B) return success, seg_positions, seg_dt def _map_curobo_to_sim( self, full_positions: torch.Tensor, curobo_joint_names: list[str], - profile: CuroboRobotProfileCfg, + backend: "_CuroboBackend", ) -> torch.Tensor: """Map a full cuRobo trajectory to simulator control-part joint order.""" - sim_to_curobo = profile.sim_to_curobo_joint_names + sim_to_curobo = backend.profile.sim_to_curobo_joint_names cols: list[int] = [] - for sim_name in sim_to_curobo: + for sim_name in backend.sim_joint_names: cu_name = sim_to_curobo[sim_name] if cu_name not in curobo_joint_names: logger.log_error( @@ -622,27 +913,54 @@ def _map_curobo_to_sim( return full_positions[..., cols].to(dtype=torch.float32) def _extract_dt( - self, v2_result: "Any", traj: "Any", max_len: int, B: int + self, + traj: "Any", + last_tstep: torch.Tensor, + max_len: int, + B: int, ) -> torch.Tensor: - """Derive ``(B, max_len)`` per-sample dt from the V2 trajectory.""" + """Derive ``(B, max_len)`` per-point deltas from a V2 trajectory. + + cuRobo V2 uses a scalar ``dt`` per batch/seed for interpolated + trajectories. EmbodiChain represents deltas at each trajectory point, + with a zero first point and one interval per following point. + """ raw_dt = getattr(traj, "dt", None) - dt = None + dt: torch.Tensor | None = None if isinstance(raw_dt, torch.Tensor): if raw_dt.dim() == 1: dt = raw_dt.unsqueeze(0).expand(B, -1) elif raw_dt.dim() == 2: dt = raw_dt if dt is None: - return torch.full( - (B, max_len), + dt = torch.full( + (B, 1), float(self.cfg.interpolation_dt), device=self.device, dtype=torch.float32, ) - T = dt.shape[-1] + if dt.shape[0] == 1 and B > 1: + dt = dt.expand(B, -1) + if dt.shape[0] != B: + logger.log_error( + f"cuRobo trajectory dt batch {dt.shape[0]} does not match {B}.", + ValueError, + ) + out = torch.zeros(B, max_len, device=self.device, dtype=torch.float32) - length = min(T, max_len) - out[:, :length] = dt[:, :length].to(self.device) + if dt.shape[-1] == 1: + interval = dt[:, 0].to(self.device, dtype=torch.float32) + for b in range(B): + length = min(int(last_tstep[b].item()) + 1, max_len) + if length > 1: + out[b, 1:length] = interval[b] + return out + + # Preserve an explicitly per-point delta sequence supplied by a V2 + # result or a compatible future API. It already includes the first + # point's zero delta in EmbodiChain's convention. + length = min(dt.shape[-1], max_len) + out[:, :length] = dt[:, :length].to(self.device, dtype=torch.float32) return out def _extract_total_time(self, v2_result: "Any", B: int) -> torch.Tensor: @@ -704,18 +1022,25 @@ def _to_curobo_joint_state( ) -> "Any": """Build a full cuRobo ``JointState`` from a sim-order control-part qpos. - Active joints are filled from ``current`` (reordered to cuRobo order); - non-active joints present in the cuRobo model are pinned to - ``fixed_joint_positions``. + Every active joint is filled from ``current`` in the simulator control + part order. Backend construction already rejects profiles whose active + V2 joints extend beyond that control part, which keeps collision + checking and the returned trajectory semantically aligned. """ profile = backend.profile curobo_names = list(backend.planner.joint_names) sim_to_curobo = profile.sim_to_curobo_joint_names curobo_to_sim_idx = { cu_name: idx - for idx, sim_name in enumerate(sim_to_curobo) + for idx, sim_name in enumerate(backend.sim_joint_names) for cu_name in [sim_to_curobo[sim_name]] } + if current.dim() != 2 or current.shape[1] != len(backend.sim_joint_names): + logger.log_error( + "cuRobo start/goal qpos must have shape " + f"(B, {len(backend.sim_joint_names)}), got {tuple(current.shape)}.", + ValueError, + ) B = current.shape[0] state = torch.zeros( B, len(curobo_names), device=self.device, dtype=torch.float32 @@ -723,8 +1048,12 @@ def _to_curobo_joint_state( for i, cu_name in enumerate(curobo_names): if cu_name in curobo_to_sim_idx: state[:, i] = current[:, curobo_to_sim_idx[cu_name]] - elif cu_name in profile.fixed_joint_positions: - state[:, i] = float(profile.fixed_joint_positions[cu_name]) + else: # Defensive: _validate_profile_joint_names rejects this case. + logger.log_error( + f"cuRobo active joint '{cu_name}' is not mapped to the " + "selected EmbodiChain control part.", + ValueError, + ) return self._bindings.JointState.from_position(state, joint_names=curobo_names) def _to_curobo_pose_goal( @@ -732,6 +1061,8 @@ def _to_curobo_pose_goal( ) -> "Any": """Build a cuRobo ``GoalToolPose`` from a batched world-frame pose.""" xpos = torch.as_tensor(xpos, device=self.device, dtype=torch.float32) + xpos = self._sim_world_to_curobo_base_pose(xpos, backend) + xpos = self._tcp_to_tool_pose(xpos, backend.profile) position, quaternion = _matrix_to_position_quaternion(xpos) pose = self._bindings.Pose(position=position, quaternion=quaternion) return self._bindings.GoalToolPose.from_poses( @@ -747,11 +1078,105 @@ def _to_curobo_joint_goal( qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) if qpos.dim() == 1: qpos = qpos.unsqueeze(0) - full_state = self._to_curobo_joint_state(qpos, backend) - return self._bindings.JointState.from_position( - full_state, joint_names=list(backend.planner.joint_names) + return self._to_curobo_joint_state(qpos, backend) + + def _tcp_to_tool_pose( + self, tcp_pose: torch.Tensor, profile: CuroboRobotProfileCfg + ) -> torch.Tensor: + """Convert a simulator TCP goal into the configured cuRobo tool frame.""" + if tcp_pose.dim() != 3 or tcp_pose.shape[-2:] != (4, 4): + logger.log_error( + f"Expected (B, 4, 4) TCP pose matrices, got {tuple(tcp_pose.shape)}.", + ValueError, + ) + if profile.tool_frame_to_tcp is None: + return tcp_pose + frame_to_tcp = torch.as_tensor( + profile.tool_frame_to_tcp, + dtype=torch.float32, + device=self.device, + ) + if frame_to_tcp.shape != (4, 4): + logger.log_error( + "tool_frame_to_tcp must be a homogeneous (4, 4) transform, " + f"got {tuple(frame_to_tcp.shape)}.", + ValueError, + ) + tool_to_frame = torch.linalg.inv(frame_to_tcp) + return tcp_pose @ tool_to_frame + + def _sim_world_to_curobo_base_pose( + self, world_pose: torch.Tensor, backend: "_CuroboBackend" + ) -> torch.Tensor: + """Express simulator-world poses in the loaded cuRobo base frame. + + EmbodiChain pose targets and dynamic obstacle poses are world poses, + while a cuRobo robot profile/world is rooted at the profile's base + link. The live simulator base pose accounts for arena offsets and + mobile bases; ``sim_base_to_curobo_base`` accounts for any fixed frame + convention difference between the two robot descriptions. + """ + if world_pose.dim() != 3 or world_pose.shape[-2:] != (4, 4): + logger.log_error( + f"Expected (B, 4, 4) simulator-world pose matrices, got " + f"{tuple(world_pose.shape)}.", + ValueError, + ) + batch_size = world_pose.shape[0] + sim_base_pose = self._get_sim_base_pose(backend, batch_size) + profile_transform = backend.profile.sim_base_to_curobo_base + if profile_transform is None: + sim_base_to_curobo = torch.eye( + 4, dtype=torch.float32, device=self.device + ).expand(batch_size, -1, -1) + else: + sim_base_to_curobo = torch.as_tensor( + profile_transform, dtype=torch.float32, device=self.device + ) + if sim_base_to_curobo.shape != (4, 4): + logger.log_error( + "sim_base_to_curobo_base must be a homogeneous (4, 4) " + f"transform, got {tuple(sim_base_to_curobo.shape)}.", + ValueError, + ) + sim_base_to_curobo = sim_base_to_curobo.expand(batch_size, -1, -1) + return torch.bmm( + sim_base_to_curobo, + torch.bmm(torch.linalg.inv(sim_base_pose), world_pose), ) + def _get_sim_base_pose( + self, backend: "_CuroboBackend", batch_size: int + ) -> torch.Tensor: + """Return ``(B, 4, 4)`` world poses of a control part's solver base.""" + control_part = backend.control_part + root_link_name = backend.profile.sim_base_link_name + if root_link_name is None: + solver = self.robot.get_solver(name=control_part) + root_link_name = getattr(solver, "root_link_name", None) + if root_link_name is None: + logger.log_error( + f"Control part '{control_part}' needs either a solver with " + "root_link_name or CuroboRobotProfileCfg.sim_base_link_name " + "for cuRobo world-frame conversion.", + ValueError, + ) + base_pose = self.robot.get_link_pose( + link_name=root_link_name, + env_ids=list(range(batch_size)), + to_matrix=True, + ) + base_pose = torch.as_tensor(base_pose, dtype=torch.float32, device=self.device) + if base_pose.dim() == 2: + base_pose = base_pose.unsqueeze(0) + if base_pose.shape != (batch_size, 4, 4): + logger.log_error( + f"Simulator base pose for '{control_part}' must have shape " + f"({batch_size}, 4, 4), got {tuple(base_pose.shape)}.", + ValueError, + ) + return base_pose + # ------------------------------------------------------------------ # Collision world + lifecycle # ------------------------------------------------------------------ @@ -767,7 +1192,8 @@ def update_dynamic_obstacles( poses: Mapping of obstacle name -> ``(B, 4, 4)`` world pose. ``None`` is a no-op. backend: Specific backend to update. If ``None``, updates all cached - backends. + backends. In ``multi_env`` mode, cached backends must share one + batch size; otherwise pass the intended backend explicitly. """ if poses is None: return @@ -775,20 +1201,59 @@ def update_dynamic_obstacles( backends = ( [backend] if backend is not None else list(self._backend_cache.values()) ) + if backend is None and self.cfg.world.multi_env: + batch_sizes = {cached_backend.batch_size for cached_backend in backends} + if len(batch_sizes) > 1: + logger.log_error( + "Cannot update all cached multi-env cuRobo backends with " + "different cached batch sizes. Pass the intended backend " + "explicitly.", + ValueError, + ) for name, pose_tensor in poses.items(): pose_tensor = torch.as_tensor( pose_tensor, device=self.device, dtype=torch.float32 ) - position, quaternion = _matrix_to_position_quaternion(pose_tensor) - B = pose_tensor.shape[0] - for b in range(B): + for be in backends: + curobo_pose = self._sim_world_to_curobo_base_pose(pose_tensor, be) + self._update_backend_obstacle(name, curobo_pose, be) + + def _update_backend_obstacle( + self, name: str, pose_tensor: torch.Tensor, backend: "_CuroboBackend" + ) -> None: + """Apply one named obstacle pose tensor under the backend's world policy.""" + if self.cfg.world.multi_env: + if pose_tensor.shape[0] != backend.batch_size: + logger.log_error( + f"dynamic obstacle '{name}' has batch {pose_tensor.shape[0]}, " + f"but this multi-env cuRobo backend expects {backend.batch_size}.", + ValueError, + ) + positions, quaternions = _matrix_to_position_quaternion(pose_tensor) + for env_idx in range(backend.batch_size): pose = self._bindings.Pose( - position=position[b], quaternion=quaternion[b] + position=positions[env_idx], quaternion=quaternions[env_idx] ) - for be in backends: - be.planner.scene_collision_checker.update_obstacle_pose( - name, pose, env_idx=b - ) + backend.planner.scene_collision_checker.update_obstacle_pose( + name, pose, env_idx=env_idx + ) + return + + # A shared world has one collision environment, so a batched input is + # only meaningful if every environment supplied the same world pose. + if pose_tensor.shape[0] > 1 and not torch.allclose( + pose_tensor, pose_tensor[:1].expand_as(pose_tensor) + ): + logger.log_error( + f"dynamic obstacle '{name}' has different poses across a " + "shared cuRobo world. Enable world.multi_env for per-env worlds.", + ValueError, + ) + position, quaternion = _matrix_to_position_quaternion(pose_tensor[:1]) + pose = self._bindings.Pose(position=position[0], quaternion=quaternion[0]) + backend.planner.scene_collision_checker.update_obstacle_pose( + name, pose, env_idx=0 + ) def close(self) -> None: """Destroy every cached cuRobo planner and clear the cache.""" @@ -816,6 +1281,8 @@ class _CuroboBackend: """Internal bundle of a cached V2 planner and its EmbodiChain-side metadata.""" planner: "Any" - tool_frame: str | None + control_part: str + sim_joint_names: list[str] + tool_frame: str profile: CuroboRobotProfileCfg batch_size: int diff --git a/embodichain/lab/sim/planners/toppra_planner.py b/embodichain/lab/sim/planners/toppra_planner.py index bafc5dda..f328a774 100644 --- a/embodichain/lab/sim/planners/toppra_planner.py +++ b/embodichain/lab/sim/planners/toppra_planner.py @@ -284,6 +284,10 @@ class ToppraPlanOptions(PlanOptions): class ToppraPlanner(BasePlanner): + """Time-optimal joint-space planner backed by TOPPRA.""" + + supports_joint_move = True + def __init__(self, cfg: ToppraPlannerCfg): r"""Initialize the TOPPRA trajectory planner. diff --git a/examples/sim/motion_gen/curobo_motion_gen.py b/examples/sim/motion_gen/curobo_motion_gen.py deleted file mode 100644 index c9dcf445..00000000 --- a/examples/sim/motion_gen/curobo_motion_gen.py +++ /dev/null @@ -1,225 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# 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. -# ---------------------------------------------------------------------------- - -"""cuRobo V2 collision-aware motion-generation demo. - -Builds a single-arm Franka Panda with a static cuboid obstacle (mirrored in -both DexSim and the cuRobo collision world), plans a collision-free -end-effector move through the EmbodiChain ``MotionGenerator`` API, and replays -the returned full-DoF trajectory in the simulator. - -Requirements: an NVIDIA CUDA device and cuRobo V2 installed with matching -CUDA/PyTorch extras. See: -https://nvlabs.github.io/curobo/latest/getting-started/installation.html -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import torch - -from embodichain import data as _data -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg -from embodichain.lab.sim.objects import RigidObjectCfg -from embodichain.lab.sim.robots import FrankaPandaCfg -from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.sim.planners import ( - MotionGenCfg, - MotionGenOptions, - MotionGenerator, - MoveType, - PlanState, -) -from embodichain.lab.sim.planners.curobo_planner import ( - CuroboPlanOptions, - CuroboPlannerCfg, - CuroboRobotProfileCfg, - CuroboWorldCfg, -) - -ROBOT_UID = "curobo_franka" -CONTROL_PART = "arm" -DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] -DEMO_BLOCK_POS = [0.45, 0.0, 0.18] -CUROBO_INSTALL_URL = ( - "https://nvlabs.github.io/curobo/latest/getting-started/installation.html" -) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="cuRobo V2 motion-gen demo") - parser.add_argument( - "--headless", - action="store_true", - help="Run without opening the viewer window.", - ) - parser.add_argument( - "--step-repeat", - type=int, - default=4, - help="Simulation updates per planned waypoint during playback.", - ) - parser.add_argument( - "--hold-steps", - type=int, - default=20, - help="Simulation updates to hold before and after playback.", - ) - parser.add_argument( - "--no-warmup", - action="store_true", - help="Skip cuRobo planner warmup.", - ) - return parser.parse_args() - - -def _check_runtime() -> None: - """Fail fast with actionable guidance when CUDA/cuRobo are unavailable.""" - if not torch.cuda.is_available(): - raise RuntimeError( - "cuRobo V2 requires a CUDA-capable NVIDIA GPU, but none is " - "available. This demo cannot run on CPU." - ) - try: - import curobo # noqa: F401 - except ImportError as exc: - raise ImportError( - "cuRobo V2 is not installed. Install it with NVIDIA's CUDA-matched " - "extras, e.g. `pip install .[cu12]` or `pip install .[cu13]` " - f"(also `.[cu12-torch]` / `.[cu13-torch]`). See {CUROBO_INSTALL_URL}." - ) from exc - - -def _demo_world_path() -> str: - return str( - Path(_data.__file__).parent - / "assets" - / "curobo" - / "collision_franka_demo.yml" - ) - - -def _franka_profile() -> CuroboRobotProfileCfg: - """Explicit sim->cuRobo joint mapping; no index order is assumed.""" - sim_to_curobo = {f"fr3_joint{i}": f"panda_joint{i}" for i in range(1, 8)} - return CuroboRobotProfileCfg( - robot_config_path="franka.yml", - sim_to_curobo_joint_names=sim_to_curobo, - fixed_joint_positions={ - "panda_finger_joint1": 0.04, - "panda_finger_joint2": 0.04, - }, - base_link_name="panda_link0", - tool_frame_name="panda_hand", - ) - - -def _build_scene(args: argparse.Namespace): - sim = SimulationManager( - SimulationManagerCfg( - headless=args.headless, - sim_device="cuda", - num_envs=1, - arena_space=2.0, - ) - ) - robot = sim.add_robot( - cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) - ) - # Mirror the cuRobo cuboid in DexSim so planner and simulator agree. - sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="demo_block", - shape=CubeCfg(size=DEMO_BLOCK_DIMS), - attrs=RigidBodyAttributesCfg(), - body_type="kinematic", - init_pos=DEMO_BLOCK_POS, - init_rot=[0.0, 0.0, 0.0], - ) - ) - return sim, robot - - -def _target_beyond_block(robot) -> torch.Tensor: - """An end-effector target that requires routing around the cuboid.""" - qpos = robot.get_qpos(name=CONTROL_PART) - fk = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True) - target = fk[0].clone() - target[:3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) - return target - - -def _play(sim, robot, trajectory: torch.Tensor, step_repeat: int) -> None: - all_ids = list(range(robot.dof)) - for w in range(trajectory.shape[1]): - robot.set_qpos(qpos=trajectory[:, w], joint_ids=all_ids) - sim.update(step=step_repeat) - - -def main() -> None: - args = parse_args() - _check_runtime() - - sim, robot = _build_scene(args) - if not args.headless: - sim.open_window() - sim.update(step=args.hold_steps) - - cfg = CuroboPlannerCfg( - robot_uid=ROBOT_UID, - robot_profiles={CONTROL_PART: _franka_profile()}, - world=CuroboWorldCfg(world_config_path=_demo_world_path()), - warmup=not args.no_warmup, - ) - mg = MotionGenerator(MotionGenCfg(planner_cfg=cfg)) - - start_qpos = robot.get_qpos(name=CONTROL_PART) - target = _target_beyond_block(robot) - - result = mg.generate( - [PlanState.from_xpos(target.unsqueeze(0))], - MotionGenOptions( - start_qpos=start_qpos, - control_part=CONTROL_PART, - plan_opts=CuroboPlanOptions(control_part=CONTROL_PART), - ), - ) - - print(f"cuRobo success: {bool(result.success.item())}") - print(f"positions shape: {tuple(result.positions.shape)}") - print(f"duration: {float(result.duration[0]):.3f}s") - - if not bool(result.success.item()): - sim.destroy() - raise RuntimeError("cuRobo failed to find a collision-free trajectory.") - - # Replay the full-DoF trajectory. - _play(sim, robot, result.positions, step_repeat=max(args.step_repeat, 1)) - sim.update(step=args.hold_steps) - - final_q = result.positions[0:1, -1, :].to(robot.device) - fk = robot.compute_fk(qpos=final_q, name=CONTROL_PART, to_matrix=True) - err = float(torch.norm(fk[0, :3, 3] - target[:3, 3])) - print(f"final Cartesian position error: {err:.4f} m") - - sim.destroy() - - -if __name__ == "__main__": - main() diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py new file mode 100644 index 00000000..8fda849d --- /dev/null +++ b/examples/sim/planners/curobo_planner.py @@ -0,0 +1,359 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""cuRobo V2 collision-aware planning through the atomic-action interface. + +The demo creates a single Franka Panda and a static cuboid that is represented +both in DexSim and in the cuRobo collision world. It then executes a +``MoveEndEffector`` action through :class:`AtomicActionEngine`, replays the +returned full-robot-DoF trajectory, and reports the final TCP error. + +Run from the repository root:: + + python examples/sim/planners/curobo_planner.py --headless + +Requirements: an NVIDIA CUDA device and cuRobo V2 installed with CUDA/PyTorch +extras compatible with the active environment. Installation instructions: +https://nvlabs.github.io/curobo/latest/getting-started/installation.html +""" + +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +import torch + +from embodichain import data as _data +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + MoveEndEffector, + MoveEndEffectorCfg, +) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.objects import RigidObjectCfg, Robot +from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator +from embodichain.lab.sim.planners.curobo_planner import ( + CuroboPlannerCfg, + CuroboRobotProfileCfg, + CuroboWorldCfg, +) +from embodichain.lab.sim.robots import FrankaPandaCfg +from embodichain.lab.sim.shapes import CubeCfg + +__all__ = ["main"] + + +ROBOT_UID = "curobo_franka" +CONTROL_PART = "arm" +DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] +DEMO_BLOCK_POS = [0.45, 0.0, 0.18] +DEFAULT_RECORD_FPS = 20 +DEFAULT_RECORD_MAX_MEMORY = 2048 +DEFAULT_RECORD_LOOK_AT = ( + (1.8, -1.8, 1.35), + (0.35, 0.10, 0.40), + (0.0, 0.0, 1.0), +) +CUROBO_INSTALL_URL = ( + "https://nvlabs.github.io/curobo/latest/getting-started/installation.html" +) + + +def parse_args() -> argparse.Namespace: + """Parse the interactive/headless playback and recording controls.""" + parser = argparse.ArgumentParser( + description="Run cuRobo V2 through EmbodiChain AtomicActionEngine." + ) + parser.add_argument( + "--headless", + action="store_true", + help="Run without opening the simulation viewer.", + ) + parser.add_argument( + "--step-repeat", + type=int, + default=4, + help="Simulation updates for each planned trajectory waypoint.", + ) + parser.add_argument( + "--hold-steps", + type=int, + default=20, + help="Simulation updates to hold before and after trajectory playback.", + ) + parser.add_argument( + "--no-warmup", + action="store_true", + help="Skip cuRobo planner warmup (useful for iteration/debugging).", + ) + parser.add_argument( + "--cuda-graph", + action="store_true", + help=( + "Enable cuRobo CUDA graphs. Disabled by default because graph capture " + "can conflict with DexSim's GPU physics stream." + ), + ) + parser.add_argument( + "--record-fps", + type=int, + default=DEFAULT_RECORD_FPS, + help="Output video FPS for automatic headless recording.", + ) + parser.add_argument( + "--record-save-path", + type=str, + default=None, + help="Optional MP4 output path for headless recording.", + ) + parser.add_argument( + "--disable-record", + action="store_true", + help="Disable automatic offscreen recording in headless mode.", + ) + return parser.parse_args() + + +def _check_runtime() -> None: + """Raise clear errors before allocating the CUDA simulation scene.""" + if not torch.cuda.is_available(): + raise RuntimeError( + "cuRobo V2 requires a CUDA-capable NVIDIA GPU, but CUDA is not " + "available. This demo cannot run on CPU." + ) + try: + import curobo # noqa: F401 + except ImportError as exc: + raise ImportError( + "cuRobo V2 is not installed. Install NVIDIA's CUDA-matched extras, " + "for example `pip install .[cu12]` or `pip install .[cu13]` " + f"(see {CUROBO_INSTALL_URL})." + ) from exc + + +def _demo_world_path() -> str: + """Return the static collision scene shared by the simulator and cuRobo.""" + return str( + Path(_data.__file__).parent / "assets" / "curobo" / "collision_franka_demo.yml" + ) + + +def _franka_profile() -> CuroboRobotProfileCfg: + """Build the explicit Franka joint and TCP-frame mapping for cuRobo.""" + return CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={ + f"fr3_joint{i}": f"panda_joint{i}" for i in range(1, 8) + }, + base_link_name="panda_link0", + tool_frame_name="panda_hand", + tool_frame_to_tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.1034], + [0.0, 0.0, 0.0, 1.0], + ], + ) + + +def _build_scene(headless: bool) -> tuple[SimulationManager, Robot]: + """Create the one-environment Franka scene with its shared cuboid.""" + sim = SimulationManager( + SimulationManagerCfg( + headless=headless, + sim_device="cuda", + num_envs=1, + arena_space=2.0, + ) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) + ) + # Keep this geometry synchronized with collision_franka_demo.yml. + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="demo_block", + shape=CubeCfg(size=DEMO_BLOCK_DIMS), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=DEMO_BLOCK_POS, + init_rot=[0.0, 0.0, 0.0], + ) + ) + return sim, robot + + +def _start_headless_recording(sim: SimulationManager, args: argparse.Namespace) -> bool: + """Start the fixed-pose offscreen recorder for a headless demo run.""" + if not args.headless or args.disable_record: + return False + if not sim.start_window_record( + save_path=args.record_save_path, + fps=args.record_fps, + max_memory=DEFAULT_RECORD_MAX_MEMORY, + video_prefix="curobo_planner_headless", + look_at=DEFAULT_RECORD_LOOK_AT, + use_sim_time=True, + ): + raise RuntimeError("Failed to start cuRobo demo headless recording.") + print("[INFO]: Headless offscreen recording enabled.") + print( + "[INFO]: The MP4 output path is reported by " + "`SimulationManager.start_window_record()`." + ) + return True + + +def _target_beyond_block(robot: Robot) -> torch.Tensor: + """Return a reachable TCP target whose route must pass around the cuboid.""" + qpos = robot.get_qpos(name=CONTROL_PART) + target = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True)[0].clone() + target[:3, 3] = torch.tensor([0.55, 0.30, 0.45], device=robot.device) + return target + + +def _replay_full_dof_trajectory( + sim: SimulationManager, + robot: Robot, + trajectory: torch.Tensor, + *, + step_repeat: int, +) -> None: + """Replay the engine's ``(B, N, robot.dof)`` trajectory in DexSim.""" + if trajectory.dim() != 3 or trajectory.shape[0] != 1: + raise ValueError( + "This single-environment demo expected a (1, N, robot.dof) " + f"trajectory, got {tuple(trajectory.shape)}." + ) + if trajectory.shape[-1] != robot.dof: + raise ValueError( + "AtomicActionEngine must return full-robot DoF positions; got " + f"{trajectory.shape[-1]} DoF for a {robot.dof}-DoF robot." + ) + + all_joint_ids = list(range(robot.dof)) + for waypoint_idx in range(trajectory.shape[1]): + waypoint = trajectory[:, waypoint_idx] + # Synchronize current state as well as the drive target. Updating a + # target alone makes the viewer show controller lag instead of the + # collision-free cuRobo waypoint being replayed. + robot.set_qpos( + qpos=waypoint, + joint_ids=all_joint_ids, + target=False, + ) + robot.set_qpos( + qpos=waypoint, + joint_ids=all_joint_ids, + target=True, + ) + sim.update(step=step_repeat) + + +def _final_tcp_error(robot: Robot, target: torch.Tensor) -> float: + """Return the Cartesian position error of the simulator's final TCP pose.""" + final_qpos = robot.get_qpos(name=CONTROL_PART) + final_pose = robot.compute_fk( + qpos=final_qpos, + name=CONTROL_PART, + to_matrix=True, + ) + return float(torch.linalg.vector_norm(final_pose[0, :3, 3] - target[:3, 3])) + + +def main() -> None: + """Plan and replay one collision-aware atomic end-effector action.""" + args = parse_args() + if args.step_repeat < 1: + raise ValueError("--step-repeat must be at least 1.") + if args.hold_steps < 0: + raise ValueError("--hold-steps must be non-negative.") + if args.record_fps < 1: + raise ValueError("--record-fps must be at least 1.") + _check_runtime() + + sim: SimulationManager | None = None + try: + sim, robot = _build_scene(args.headless) + if not args.headless: + sim.open_window() + _start_headless_recording(sim, args) + if args.hold_steps: + sim.update(step=args.hold_steps) + + motion_generator = MotionGenerator( + MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg(world_config_path=_demo_world_path()), + warmup=not args.no_warmup, + use_cuda_graph=args.cuda_graph, + ) + ) + ) + engine = AtomicActionEngine(motion_generator) + engine.register( + MoveEndEffector( + motion_generator, + MoveEndEffectorCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part=CONTROL_PART, + sample_interval=80, + ), + ), + name="move_end_effector", + ) + + target = _target_beyond_block(robot) + plan_start = time.perf_counter() + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + planning_duration = time.perf_counter() - plan_start + + print(f"cuRobo atomic-action success: {bool(success.item())}") + print(f"full-DoF trajectory shape: {tuple(trajectory.shape)}") + print(f"atomic-action planning duration: {planning_duration:.3f} s") + + if not bool(success.item()): + raise RuntimeError("cuRobo failed to find a collision-free trajectory.") + + _replay_full_dof_trajectory( + sim, + robot, + trajectory, + step_repeat=args.step_repeat, + ) + if args.hold_steps: + sim.update(step=args.hold_steps) + print(f"final TCP position error: {_final_tcp_error(robot, target):.4f} m") + finally: + if sim is not None: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + sim.destroy() + SimulationManager.flush_cleanup_queue() + + +if __name__ == "__main__": + main() diff --git a/tests/docs/test_curobo_planner_docs.py b/tests/docs/test_curobo_planner_docs.py new file mode 100644 index 00000000..ff1f6a97 --- /dev/null +++ b/tests/docs/test_curobo_planner_docs.py @@ -0,0 +1,42 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# 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. +# ---------------------------------------------------------------------------- + +"""Source-level coverage for the public cuRobo planner documentation.""" + +from __future__ import annotations + +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_PLANNER_DOCS = _REPO_ROOT / "docs" / "source" / "overview" / "sim" / "planners" + + +def test_curobo_planner_docs_are_linked_and_scoped() -> None: + """Keep the optional V2 backend discoverable without overstating support.""" + index = (_PLANNER_DOCS / "index.rst").read_text(encoding="utf-8") + page = (_PLANNER_DOCS / "curobo_planner.md").read_text(encoding="utf-8") + + assert "curobo_planner.md" in index + assert "CuroboPlannerCfg" in page + assert 'planner_type="curobo"' in page + assert "cuRobo V2" in page + assert "attached-object" in page + assert "tool_frame_to_tcp" in page + assert "sim_base_to_curobo_base" in page + assert "multi_env=True" in page + assert "lock_joints" in page + assert "--record-save-path" in page + assert "examples/sim/planners/curobo_planner.py" in page diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 1a88cbf6..d315413b 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -135,6 +135,7 @@ def _make_curobo_mock_motion_generator(result_positions, success=None): planner.cfg.planner_type = "curobo" planner.preinterpolate_targets = False planner.preserve_plan_samples = True + planner.supports_joint_move = True mg.planner = planner B = result_positions.shape[0] if success is None: diff --git a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py index 15165960..e156cb8f 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -65,10 +65,7 @@ def _demo_world_path() -> str: return str( - Path(_data.__file__).parent - / "assets" - / "curobo" - / "collision_franka_demo.yml" + Path(_data.__file__).parent / "assets" / "curobo" / "collision_franka_demo.yml" ) @@ -77,12 +74,14 @@ def _franka_profile() -> CuroboRobotProfileCfg: return CuroboRobotProfileCfg( robot_config_path="franka.yml", sim_to_curobo_joint_names=sim_to_curobo, - fixed_joint_positions={ - "panda_finger_joint1": 0.04, - "panda_finger_joint2": 0.04, - }, base_link_name="panda_link0", tool_frame_name="panda_hand", + tool_frame_to_tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.1034], + [0.0, 0.0, 0.0, 1.0], + ], ) @@ -109,6 +108,8 @@ def _make_franka_curobo_engine(): robot_uid=ROBOT_UID, robot_profiles={CONTROL_PART: _franka_profile()}, world=CuroboWorldCfg(world_config_path=_demo_world_path()), + warmup=False, + use_cuda_graph=False, ) ) ) @@ -133,14 +134,24 @@ def _reachable_target_beyond_demo_block(robot) -> torch.Tensor: qpos = robot.get_qpos(name=CONTROL_PART) fk = robot.compute_fk(qpos=qpos, name=CONTROL_PART, to_matrix=True) target = fk[0].clone() - target[:3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) + target[:3, 3] = torch.tensor([0.55, 0.30, 0.45], device=robot.device) return target def _play_trajectory(sim, robot, trajectory: torch.Tensor, step_repeat: int = 1): - arm_ids = robot.get_joint_ids(name=CONTROL_PART) + """Replay every waypoint with matching state and drive targets. + + ``target=True`` alone only updates the articulation drive target. The + physics step can therefore still be catching up with the previous sample + when the next one is supplied. Set the current state first so the replay + validates the planned configuration rather than the drive controller's + tracking transient. + """ + all_joint_ids = list(range(robot.dof)) for w in range(trajectory.shape[1]): - robot.set_qpos(qpos=trajectory[:, w], joint_ids=list(range(robot.dof))) + waypoint = trajectory[:, w] + robot.set_qpos(qpos=waypoint, joint_ids=all_joint_ids, target=False) + robot.set_qpos(qpos=waypoint, joint_ids=all_joint_ids, target=True) sim.update(step=step_repeat) diff --git a/tests/sim/atomic_actions/test_trajectory_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py index dd552700..bcebd90a 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -20,7 +20,7 @@ import torch import pytest -from unittest.mock import Mock +from unittest.mock import Mock, patch from embodichain.lab.sim.atomic_actions.trajectory import TrajectoryBuilder from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType @@ -45,6 +45,7 @@ def compute_ik(pose=None, name=None, joint_seed=None, **kw): # TOPPRA allows MotionGenerator pre-interpolation; neural/curobo do not. planner.preinterpolate_targets = planner_type == "toppra" planner.preserve_plan_samples = planner_type == "curobo" + planner.supports_joint_move = planner_type in {"toppra", "curobo"} mg.planner = planner return mg @@ -110,16 +111,21 @@ def test_ik_interp_path_unchanged(self): [PlanState(xpos=torch.eye(4), move_type=MoveType.EEF_MOVE)] for _ in range(2) ] - ok, traj = builder.plan_arm_traj( - target_states_list, - start_qpos, - 10, - control_part="arm", - arm_dof=6, - cfg=cfg, - ) + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + return_value=torch.zeros(2, 10, 6), + ) as interpolate: + ok, traj = builder.plan_arm_traj( + target_states_list, + start_qpos, + 10, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) assert ok.all().item() assert traj.shape[0] == 2 + interpolate.assert_called_once() class TestCuroboBuilderDispatch: @@ -239,3 +245,82 @@ def test_failed_row_holds_start_qpos(self): assert success.tolist() == [True, False] # Failed env held at its start qpos across all samples. assert torch.allclose(trajectory[1], start[1].unsqueeze(0).repeat(5, 1)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def test_indexless_cuda_robot_accepts_indexed_curobo_result(self): + positions = torch.zeros(2, 5, 6, device="cuda:0") + mg = _mock_curobo_motion_generator(result_positions=positions) + mg.robot.device = torch.device("cuda") + mg.device = torch.device("cuda") + builder = TrajectoryBuilder(mg) + + success, trajectory = builder.plan_arm_traj( + _pose_targets_for_two_envs(), + torch.zeros(2, 6, device="cuda:0"), + n_waypoints=10, + control_part="arm", + arm_dof=6, + cfg=ActionCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part="arm", + ), + ) + + assert success.tolist() == [True, True] + assert trajectory.device == torch.device("cuda:0") + + +class TestJointMotionCapabilities: + def test_joint_capable_backend_delegates_through_motion_generator(self): + """TOPPRA retains the established motion-generator joint route.""" + mg = _mock_mg(num_envs=2, arm_dof=6, planner_type="toppra") + mg.generate.return_value = PlanResult( + success=torch.ones(2, dtype=torch.bool), + positions=torch.zeros(2, 8, 6), + ) + builder = TrajectoryBuilder(mg) + cfg = ActionCfg( + motion_source="motion_gen", planner_type="toppra", control_part="arm" + ) + + success, trajectory = builder.plan_joint_motion( + torch.zeros(2, 6), + torch.ones(2, 6), + n_waypoints=8, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) + + assert success.tolist() == [True, True] + assert trajectory.shape == (2, 8, 6) + assert mg.generate.call_args.args[0][0].move_type is MoveType.JOINT_MOVE + + def test_neural_joint_motion_falls_back_to_local_interpolation(self): + """Neural is Cartesian-only, so joint phases must not call generate.""" + mg = _mock_mg(num_envs=2, arm_dof=6, planner_type="neural") + builder = TrajectoryBuilder(mg) + start = torch.zeros(2, 6) + target = torch.ones(2, 6) + cfg = ActionCfg( + motion_source="motion_gen", planner_type="neural", control_part="arm" + ) + + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + return_value=torch.zeros(2, 8, 6), + ) as interpolate: + success, trajectory = builder.plan_joint_motion( + start, + target, + n_waypoints=8, + control_part="arm", + arm_dof=6, + cfg=cfg, + ) + + assert success.tolist() == [True, True] + assert trajectory.shape == (2, 8, 6) + interpolate.assert_called_once() + mg.generate.assert_not_called() diff --git a/tests/sim/planners/test_curobo_integration.py b/tests/sim/planners/test_curobo_integration.py index 4ed2514f..41af4f06 100644 --- a/tests/sim/planners/test_curobo_integration.py +++ b/tests/sim/planners/test_curobo_integration.py @@ -50,6 +50,7 @@ ) from embodichain.lab.sim.planners.curobo_planner import ( # noqa: E402 CuroboPlanOptions, + CuroboPlanner, CuroboPlannerCfg, CuroboRobotProfileCfg, CuroboWorldCfg, @@ -59,6 +60,10 @@ CONTROL_PART = "arm" DEMO_BLOCK_DIMS = [0.18, 0.40, 0.36] DEMO_BLOCK_POS = [0.45, 0.0, 0.18] +# Small displacement from Panda's neutral ready pose. It is deliberately far +# from the joint limits so this smoke test exercises cuRobo planning rather +# than limit handling. +JOINT_1_TARGET_DELTA_RAD = 0.12 def _demo_world_path() -> str: @@ -73,18 +78,23 @@ def _franka_profile() -> CuroboRobotProfileCfg: return CuroboRobotProfileCfg( robot_config_path="franka.yml", sim_to_curobo_joint_names=sim_to_curobo, - fixed_joint_positions={ - "panda_finger_joint1": 0.04, - "panda_finger_joint2": 0.04, - }, base_link_name="panda_link0", tool_frame_name="panda_hand", + # DexSim targets the Panda TCP, while the stock cuRobo profile uses + # panda_hand. This fixed transform converts the requested TCP pose + # back into the cuRobo hand frame before planning. + tool_frame_to_tcp=[ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.1034], + [0.0, 0.0, 0.0, 1.0], + ], ) -def _make_sim_robot(): +def _make_sim_robot(num_envs: int = 1): sim = SimulationManager( - SimulationManagerCfg(headless=True, sim_device="cuda", num_envs=1) + SimulationManagerCfg(headless=True, sim_device="cuda", num_envs=num_envs) ) robot = sim.add_robot( cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) @@ -111,6 +121,10 @@ def test_curobo_v2_plans_around_a_static_cuboid(): robot_uid=ROBOT_UID, robot_profiles={CONTROL_PART: _franka_profile()}, world=CuroboWorldCfg(world_config_path=_demo_world_path()), + # The real smoke test validates planner calls, not CUDA-graph capture. + # Skipping warmup keeps it practical on fresh CI GPU workers. + warmup=False, + use_cuda_graph=False, ) mg = MotionGenerator(MotionGenCfg(planner_cfg=cfg)) @@ -120,7 +134,7 @@ def test_curobo_v2_plans_around_a_static_cuboid(): ) # Target beyond the cuboid so the planner must route around it. target_xpos = start_xpos.clone() - target_xpos[0, :3, 3] = torch.tensor([0.55, 0.20, 0.30], device=robot.device) + target_xpos[0, :3, 3] = torch.tensor([0.55, 0.30, 0.45], device=robot.device) result = mg.generate( [PlanState.from_xpos(target_xpos)], @@ -150,3 +164,95 @@ def test_curobo_v2_plans_around_a_static_cuboid(): finally: sim.destroy() SimulationManager.flush_cleanup_queue() + + +@pytest.mark.slow +def test_curobo_v2_plans_a_joint_space_move(): + """Route a ``JOINT_MOVE`` through V2 ``plan_cspace`` on CUDA.""" + sim, robot = _make_sim_robot() + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg(world_config_path=_demo_world_path()), + warmup=False, + use_cuda_graph=False, + ) + mg = MotionGenerator(MotionGenCfg(planner_cfg=cfg)) + start_qpos = robot.get_qpos(name=CONTROL_PART) + target_qpos = start_qpos.clone() + target_qpos[:, 0] += JOINT_1_TARGET_DELTA_RAD + + result = mg.generate( + [PlanState.from_qpos(target_qpos)], + MotionGenOptions( + start_qpos=start_qpos, + control_part=CONTROL_PART, + plan_opts=CuroboPlanOptions(control_part=CONTROL_PART), + ), + ) + + assert result.success.shape == (1,) + assert bool(result.success.item()) + assert result.positions is not None + assert result.positions.shape[-1] == start_qpos.shape[-1] + assert torch.allclose(result.positions[0, 0], start_qpos[0], atol=1e-3) + assert torch.allclose(result.positions[0, -1], target_qpos[0], atol=1e-3) + assert float(result.duration[0]) > 0.0 + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() + + +@pytest.mark.slow +def test_curobo_v2_multi_env_worlds_are_independent(): + """V2 gets one static world and one dynamic obstacle update per row.""" + sim, robot = _make_sim_robot(num_envs=2) + planner = None + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + robot_profiles={CONTROL_PART: _franka_profile()}, + world=CuroboWorldCfg( + world_config_path=_demo_world_path(), + dynamic_obstacle_names=["demo_block"], + multi_env=True, + ), + warmup=False, + use_cuda_graph=False, + ) + planner = CuroboPlanner(cfg) + profile = cfg.robot_profiles[CONTROL_PART] + backend = planner._get_backend(profile, CONTROL_PART, batch_size=2) + collision_checker = backend.planner.scene_collision_checker + + assert collision_checker.num_envs == 2 + assert collision_checker.get_obstacle_names(env_idx=0) == ["demo_block"] + assert collision_checker.get_obstacle_names(env_idx=1) == ["demo_block"] + + start_qpos = robot.get_qpos(name=CONTROL_PART) + target_qpos = start_qpos.clone() + target_qpos[:, 0] += torch.tensor([0.08, -0.08], device=robot.device) + result = planner.plan( + [PlanState.from_qpos(target_qpos)], + CuroboPlanOptions(start_qpos=start_qpos, control_part=CONTROL_PART), + ) + assert result.success.tolist() == [True, True] + assert result.positions is not None + assert result.positions.shape[0] == 2 + + # Start from each live simulator base, then request different local + # offsets. This verifies the adapter writes each V2 world independently. + dynamic_poses = planner._get_sim_base_pose(backend, batch_size=2).clone() + dynamic_poses[:, 0, 3] += torch.tensor( + [0.10, -0.15], device=dynamic_poses.device + ) + planner.update_dynamic_obstacles({"demo_block": dynamic_poses}, backend) + + inv_pose = collision_checker.data.cuboids.inv_pose[:, 0, :3] + assert not torch.allclose(inv_pose[0], inv_pose[1]) + finally: + if planner is not None: + planner.close() + sim.destroy() + SimulationManager.flush_cleanup_queue() diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index 4dca2391..47357cac 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -25,6 +25,8 @@ from __future__ import annotations import importlib +from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -71,6 +73,8 @@ def test_matrix_to_position_quaternion_uses_wxyz(): position, quaternion = _matrix_to_position_quaternion(matrix) assert torch.equal(position, torch.zeros(1, 3)) assert torch.equal(quaternion, torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + assert position.is_contiguous() + assert quaternion.is_contiguous() def test_matrix_to_position_quaternion_rejects_non_4x4_batch(): @@ -111,10 +115,16 @@ def test_curobo_planner_cfg_defaults(): assert cfg.planner_type == "curobo" assert cfg.warmup is True assert cfg.max_attempts == 5 - assert cfg.use_cuda_graph is True + assert cfg.use_cuda_graph is False assert isinstance(cfg.world, CuroboWorldCfg) +def test_curobo_world_cfg_uses_v2_safe_default_collision_cache(): + cache = CuroboWorldCfg().collision_cache + + assert cache == {"cuboid": 8, "mesh": 2} + + def test_curobo_robot_profile_cfg_requires_joint_map(): cfg = CuroboRobotProfileCfg( robot_config_path="franka.yml", @@ -122,7 +132,6 @@ def test_curobo_robot_profile_cfg_requires_joint_map(): ) assert cfg.robot_config_path == "franka.yml" assert cfg.sim_to_curobo_joint_names == {"a": "b"} - assert cfg.fixed_joint_positions == {} def test_curobo_planner_class_is_lazy_import_safe(): @@ -184,8 +193,9 @@ def update_obstacle_pose(self, name, pose, env_idx=0): class _FakeKinematics: - def __init__(self, joint_names): + def __init__(self, joint_names, base_link="base"): self.joint_names = list(joint_names) + self.base_link = base_link class _FakeV2PlannerInstance: @@ -193,6 +203,7 @@ def __init__(self, bindings): self._bindings = bindings self.joint_names = list(bindings.full_joint_names) self.tool_frame = bindings.tool_frame + self.tool_frames = list(bindings.tool_frames) self.scene_collision_checker = _FakeCollisionChecker() self.kinematics = _FakeKinematics(self.joint_names) self.plan_pose_calls = [] @@ -201,6 +212,7 @@ def __init__(self, bindings): self.max_batch_size = None self.closed = False self.warmup_count = 0 + self.warmup_enable_graph = None def plan_pose(self, goal, current_state, max_attempts=5): self.plan_pose_calls.append((goal, current_state, max_attempts)) @@ -215,8 +227,9 @@ def _next_result(self): return self._bindings.results.pop(0) return self._bindings.next_result - def warmup(self): + def warmup(self, enable_graph=True): self.warmup_count += 1 + self.warmup_enable_graph = enable_graph def close(self): self.closed = True @@ -258,6 +271,8 @@ def __init__(self, bindings): self._bindings = bindings def from_position(self, position, joint_names=None): + if not isinstance(position, torch.Tensor): + raise TypeError("JointState.from_position requires a tensor position") return _FakeJointState(position=position, joint_names=joint_names) @@ -279,12 +294,18 @@ def from_poses(self, pose_dict, ordered_tool_frames=None, num_goalset=1): ) +class _FakeDeviceCfgFactory: + def __call__(self, device): + return ("fake_device_cfg", device) + + class _FakeCuroboBindings: """A minimal stand-in for the cuRobo V2 facade namespace.""" def __init__(self, full_joint_names, tool_frame="tool"): self.full_joint_names = list(full_joint_names) self.tool_frame = tool_frame + self.tool_frames = [tool_frame] self.warmup_count = 0 self.create_kwargs = None self.created_planners: list = [] @@ -295,6 +316,7 @@ def __init__(self, full_joint_names, tool_frame="tool"): self.JointState = _FakeJointStateFactory(self) self.Pose = _FakePoseFactory(self) self.GoalToolPose = _FakeGoalToolPoseFactory(self) + self.DeviceCfg = _FakeDeviceCfgFactory() self.next_result = self.make_result( position=torch.zeros(1, 1, 3, len(full_joint_names)), dt=torch.tensor([0.0, 0.025, 0.025]), @@ -325,6 +347,9 @@ def __init__(self, device="cuda", num_instances=1, dof=2): self.device = torch.device(device) self.num_instances = num_instances self.dof = dof + self.control_parts = {"arm": ["sim_a", "sim_b"]} + self._solver = SimpleNamespace(root_link_name="sim_base") + self.base_poses = torch.eye(4, device=self.device).repeat(num_instances, 1, 1) def get_qpos(self, name=None): return torch.zeros(self.num_instances, self.dof) @@ -332,6 +357,17 @@ def get_qpos(self, name=None): def get_joint_ids(self, name=None): return list(range(self.dof)) + def get_solver(self, name=None): + assert name == "arm" + return self._solver + + def get_link_pose(self, link_name, env_ids=None, to_matrix=False): + assert link_name == self._solver.root_link_name + assert to_matrix is True + if env_ids is None: + return self.base_poses + return self.base_poses[env_ids] + class _FakeSim: def __init__(self, robot): @@ -470,6 +506,196 @@ def test_unknown_control_part_raises(fake_curobo, fake_sim): ) +def test_profile_base_link_must_match_curobo_model(fake_curobo, fake_sim): + """A configured base frame must not silently disagree with the V2 model.""" + profile = _default_profile() + profile.base_link_name = "unexpected_base" + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + + with pytest.raises(ValueError, match="base_link_name"): + planner.plan( + [PlanState.from_qpos(torch.zeros(1, 2))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + +def test_active_joints_outside_control_part_are_rejected(fake_curobo, fake_sim): + """Collision planning must not pin unexecuted joints only inside cuRobo.""" + fake_curobo.full_joint_names.append("cu_gripper") + planner = _make_planner(fake_curobo, fake_sim) + + with pytest.raises(ValueError, match="outside the requested control part"): + planner.plan( + [PlanState.from_qpos(torch.zeros(1, 2))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + +def test_target_batch_must_match_planning_start_batch(fake_curobo, fake_sim): + """Direct planner calls must reject a start/target batch mismatch early.""" + fake_sim.robot = _FakeRobot(num_instances=2) + planner = _make_planner(fake_curobo, fake_sim) + + with pytest.raises(ValueError, match="target batch 2.*start batch 1"): + planner.plan( + [PlanState.from_qpos(torch.zeros(2, 2))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + +def test_pose_goal_is_expressed_in_curobo_base_frame(fake_curobo, fake_sim): + """Simulator-world targets must be rebased before cuRobo sees them.""" + fake_sim.robot.base_poses[0, 0, 3] = 10.0 + profile = _default_profile() + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + backend = planner._get_backend(profile, "arm", batch_size=1) + world_target = torch.eye(4).unsqueeze(0) + world_target[0, 0, 3] = 10.5 + + goal = planner._to_curobo_pose_goal(world_target, backend) + + assert torch.allclose( + goal.pose_dict["tool"].position.cpu(), torch.tensor([[0.5, 0.0, 0.0]]) + ) + + +def test_profile_fixed_base_transform_is_applied(fake_curobo, fake_sim): + """A profile-specific simulator-base to cuRobo-base transform is composed.""" + fake_sim.robot.base_poses[0, 0, 3] = 10.0 + profile = _default_profile() + profile.sim_base_to_curobo_base = [ + [1.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + backend = planner._get_backend(profile, "arm", batch_size=1) + world_target = torch.eye(4).unsqueeze(0) + world_target[0, 0, 3] = 10.5 + + goal = planner._to_curobo_pose_goal(world_target, backend) + + assert torch.allclose( + goal.pose_dict["tool"].position.cpu(), torch.tensor([[1.5, 0.0, 0.0]]) + ) + + +def test_profile_sim_base_link_works_without_local_solver(fake_curobo, fake_sim): + """A profile can provide the simulator base link without an IK solver.""" + fake_sim.robot.get_solver = lambda name=None: None + profile = _default_profile() + profile.sim_base_link_name = "sim_base" + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + backend = planner._get_backend(profile, "arm", batch_size=1) + + goal = planner._to_curobo_pose_goal(torch.eye(4).unsqueeze(0), backend) + + assert torch.allclose( + goal.pose_dict["tool"].position.cpu(), torch.tensor([[0.0, 0.0, 0.0]]) + ) + + +def test_dynamic_obstacle_pose_is_expressed_in_curobo_base_frame(fake_curobo, fake_sim): + """Dynamic simulator-world obstacle poses use the same base conversion.""" + fake_sim.robot.base_poses[0, 0, 3] = 10.0 + world = CuroboWorldCfg(dynamic_obstacle_names=["block"]) + planner = _make_planner(fake_curobo, fake_sim, world=world) + profile = planner.cfg.robot_profiles["arm"] + backend = planner._get_backend(profile, "arm", batch_size=1) + world_pose = torch.eye(4).unsqueeze(0) + world_pose[0, 0, 3] = 10.5 + + planner.update_dynamic_obstacles({"block": world_pose}, backend) + + _, pose, _ = backend.planner.scene_collision_checker.updates[-1] + assert torch.allclose(pose.position.cpu(), torch.tensor([0.5, 0.0, 0.0])) + + +def test_batched_pose_goals_rebase_each_simulator_arena(fake_curobo, fake_sim): + """Parallel arenas retain one common local cuRobo target per environment.""" + fake_sim.robot = _FakeRobot(num_instances=2) + fake_sim.robot.base_poses[1, 0, 3] = 10.0 + profile = _default_profile() + planner = _make_planner( + fake_curobo, + fake_sim, + profiles={"arm": profile}, + world=CuroboWorldCfg(multi_env=True), + ) + backend = planner._get_backend(profile, "arm", batch_size=2) + world_targets = torch.eye(4).unsqueeze(0).repeat(2, 1, 1) + world_targets[:, 0, 3] = torch.tensor([0.5, 10.5]) + + goal = planner._to_curobo_pose_goal(world_targets, backend) + + expected = torch.tensor([[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]]) + assert torch.allclose(goal.pose_dict["tool"].position.cpu(), expected) + + +def test_multi_env_materializes_one_scene_mapping_per_batch_row( + fake_curobo, fake_sim, tmp_path +): + """V2 needs a list of B scene mappings, not a singleton scene path.""" + scene_path = Path(tmp_path) / "world.yml" + scene_path.write_text( + "cuboid:\n" + " block:\n" + " dims: [0.1, 0.1, 0.1]\n" + " pose: [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]\n", + encoding="utf-8", + ) + fake_sim.robot = _FakeRobot(num_instances=2) + planner = _make_planner( + fake_curobo, + fake_sim, + world=CuroboWorldCfg(world_config_path=str(scene_path), multi_env=True), + ) + profile = planner.cfg.robot_profiles["arm"] + + planner._get_backend(profile, "arm", batch_size=2) + + scene_models = fake_curobo.create_kwargs["scene_model"] + assert len(scene_models) == 2 + assert scene_models[0] == scene_models[1] + assert scene_models[0] is not scene_models[1] + assert scene_models[0]["cuboid"] is not scene_models[1]["cuboid"] + + +def test_multi_env_materializes_empty_scene_for_every_batch_row(fake_curobo, fake_sim): + """An empty cached V2 collision world still needs B scene entries.""" + fake_sim.robot = _FakeRobot(num_instances=2) + planner = _make_planner( + fake_curobo, + fake_sim, + world=CuroboWorldCfg(multi_env=True), + ) + profile = planner.cfg.robot_profiles["arm"] + + planner._get_backend(profile, "arm", batch_size=2) + + assert fake_curobo.create_kwargs["scene_model"] == [{}, {}] + + +def test_dynamic_update_requires_explicit_backend_for_mixed_batch_caches( + fake_curobo, fake_sim +): + """Avoid partially updating incompatible multi-env cuRobo cache entries.""" + world = CuroboWorldCfg(multi_env=True, dynamic_obstacle_names=["block"]) + planner = _make_planner(fake_curobo, fake_sim, world=world) + profile = planner.cfg.robot_profiles["arm"] + planner._get_backend(profile, "arm", batch_size=1) + planner._get_backend(profile, "arm", batch_size=2) + + with pytest.raises(ValueError, match="different cached batch sizes"): + planner.update_dynamic_obstacles({"block": torch.eye(4).unsqueeze(0)}) + + assert all( + not backend.planner.scene_collision_checker.updates + for backend in planner._backend_cache.values() + ) + + def test_non_cuda_device_is_rejected(monkeypatch): robot = _FakeRobot(device="cpu") sim = _FakeSim(robot) @@ -506,6 +732,16 @@ def test_backend_is_cached_across_plans(fake_curobo, fake_sim): assert fake_curobo.created_planners[0].warmup_count == 1 +def test_warmup_respects_use_cuda_graph_cfg(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim, use_cuda_graph=False) + planner.plan( + [PlanState.from_qpos(torch.zeros(1, 2))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + assert fake_curobo.created_planners[0].warmup_enable_graph is False + + def test_update_dynamic_obstacles_reaches_backend(fake_curobo, fake_sim): world = CuroboWorldCfg(dynamic_obstacle_names=["block"]) planner = _make_planner(fake_curobo, fake_sim, world=world) @@ -553,3 +789,79 @@ def test_joint_move_uses_plan_cspace(fake_curobo, fake_sim): planner_inst = fake_curobo.created_planners[0] assert len(planner_inst.plan_cspace_calls) == 1 assert len(planner_inst.plan_pose_calls) == 0 + goal_state = planner_inst.plan_cspace_calls[0][0] + assert torch.allclose(goal_state.position.cpu(), target) + + +def test_none_v2_result_holds_start_qpos(fake_curobo, fake_sim): + fake_curobo.next_result = None + planner = _make_planner(fake_curobo, fake_sim) + start = torch.tensor([[0.3, -0.4]]) + + result = planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=start, control_part="arm"), + ) + + assert result.success.tolist() == [False] + assert torch.equal(result.positions.cpu(), start.unsqueeze(1)) + + +def test_scalar_v2_dt_expands_over_trajectory_intervals(fake_curobo, fake_sim): + fake_curobo.next_result = fake_curobo.make_result( + position=torch.tensor([[[[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]]]]), + dt=torch.tensor([[0.025]]), + ) + planner = _make_planner(fake_curobo, fake_sim) + + result = planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + assert torch.allclose(result.dt.cpu(), torch.tensor([[0.0, 0.025, 0.025]])) + assert torch.allclose(result.duration.cpu(), torch.tensor([0.05])) + + +def test_joint_mapping_uses_control_part_order_not_dict_insertion_order( + fake_curobo, fake_sim +): + profile = CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={"sim_b": "cu_b", "sim_a": "cu_a"}, + ) + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + backend = planner._get_backend(profile, "arm", batch_size=1) + + state = planner._to_curobo_joint_state(torch.tensor([[10.0, 20.0]]), backend) + + assert torch.allclose(state.position.cpu(), torch.tensor([[10.0, 20.0]])) + + +def test_missing_tool_frame_uses_single_backend_tool_frame(fake_curobo, fake_sim): + profile = CuroboRobotProfileCfg( + robot_config_path="franka.yml", + sim_to_curobo_joint_names={"sim_a": "cu_a", "sim_b": "cu_b"}, + tool_frame_name=None, + ) + planner = _make_planner(fake_curobo, fake_sim, profiles={"arm": profile}) + backend = planner._get_backend(profile, "arm", batch_size=1) + + goal = planner._to_curobo_pose_goal(torch.eye(4).unsqueeze(0), backend) + + assert goal.ordered_tool_frames == ["tool"] + assert list(goal.pose_dict) == ["tool"] + + +def test_backend_normalizes_indexless_cuda_device_for_warp(fake_curobo, fake_sim): + planner = _make_planner(fake_curobo, fake_sim) + + planner.plan( + [PlanState.from_xpos(torch.eye(4).unsqueeze(0))], + CuroboPlanOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + assert fake_curobo.create_kwargs["device_cfg"] == ( + "fake_device_cfg", + torch.device("cuda:0"), + )