diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 38d474e72..1bc345005 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -158,7 +158,7 @@ jobs: - uses: ./.github/actions/activate-conda - name: Install test package run: | - pip install -e ".[gensim]" \ + pip install -e ".[gensim, curobo-cu12]" \ --extra-index-url http://pyp.open3dv.site:2345/simple/ \ --trusted-host pyp.open3dv.site \ --extra-index-url https://download.blender.org/pypi/ diff --git a/.gitignore b/.gitignore index 16f43900b..69061763a 100644 --- a/.gitignore +++ b/.gitignore @@ -203,5 +203,8 @@ embodichain/VERSION # benchmark results scripts/benchmark/rl/reports/* +# Local agent execution state and isolated worktrees +.superpowers/ +.worktrees/ # Local gym project workspace /gym_project/ 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 000000000..fc5142b50 --- /dev/null +++ b/docs/source/overview/sim/planners/curobo_planner.md @@ -0,0 +1,265 @@ +# 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 + +EmbodiChain exposes cuRobo V2 as CUDA-matched optional dependencies. From the +EmbodiChain repository root, select exactly one extra: + +~~~bash +# Recommended for the normal EmbodiChain environment, where PyTorch is present. +uv pip install ".[curobo-cu12]" # CUDA 12.x +uv pip install ".[curobo-cu13]" # CUDA 13.x + +# For a fresh environment that also needs PyTorch. +uv pip install ".[curobo-cu12-torch]" # CUDA 12.x +uv pip install ".[curobo-cu13-torch]" # CUDA 13.x + +python -c "import curobo; print(curobo.__version__)" +pytest --pyargs curobo.tests +~~~ + +The extras follow [NVIDIA's official cuRobo installation +guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html) +and pin the source dependency to the cuRobo V2 `v0.8.0` release. Use a Python +3.10--3.13 environment on Linux with a supported NVIDIA GPU and driver. The +non-`torch` extras are preferred for EmbodiChain because the simulation +environment normally already provides PyTorch; the `-torch` variants delegate +the PyTorch version requirement to cuRobo. Keep cuRobo in the same Python +environment that runs the simulator. + +## Configure a control part + +The cuRobo robot model and the per-control-part profile are both auto-generated +internally - no external cuRobo robot YAML (e.g. `franka.yml`) and no +`robot_profiles` config are needed. On the first plan, the adapter fits collision +spheres to each link of the robot's URDF and writes a cuRobo V2 robot YAML (see +[Auto-generated robot YAML](curobo-auto-generated-robot-yaml). The tool frame, TCP +offset, and base link are read from the control part's IK solver, and the +simulator->cuRobo joint mapping is identity (the generated YAML reuses the +URDF's own joint names). The control part is selected at plan time through +`CuroboPlanOptions.control_part` and validated against `robot.control_parts`. + +Lock non-controlled joints (for example gripper joints) in the cuRobo 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. For example, the 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, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, +) + +planner_cfg = CuroboPlannerCfg( + robot_uid="my_franka", + planner_type="curobo", + world=CuroboWorldCfg(rigid_objects=[demo_block]), +) +motion_generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg)) +~~~ + +The physics and planner devices are independent. +`SimulationManagerCfg(sim_device="cpu")` keeps robot state, targets, and +returned trajectories on CPU, +while cuRobo still performs all model generation and planning on CUDA. By +default a CPU simulation uses PyTorch's current CUDA device; set +`CuroboPlannerCfg.cuda_device="cuda:1"` (or an integer GPU index) to select a +different planning GPU. A CPU value is rejected because cuRobo itself has no +CPU backend. + +The robot configuration must be a cuRobo V2 robot profile with collision +spheres and self-collision data; the adapter generates this from the robot's +URDF automatically. A plain URDF alone is not sufficient for collision planning +without that sphere-fitting step. + +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 +`CuroboPlannerCfg.sim_base_to_curobo_base` to the transform from the simulator +base to the cuRobo base. Collision-world poses are authored in the cuRobo +base/world frame. `tool_frame_to_tcp` (read from `solver.tcp_xpos`) converts an +EmbodiChain TCP goal into the chosen cuRobo tool frame when the solver's end link +is not itself the TCP. 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. + +`CuroboPlannerCfg.use_cuda_graph` defaults to `True`. The planner runs in the +simulator process and reuses its CUDA context; it does not launch a persistent +`spawn` worker or copy planning tensors through multiprocessing queues. Set +`use_cuda_graph=False` when lower one-time initialization cost and lower +graph-resident memory are more important than hot planning latency. + +In graph mode, planner initialization uses the same per-device +`CaptureCoordinator` as DexSim's Newton backend and synchronizes the device +before and after cuRobo warmup. EmbodiChain also forces cuRobo's PyTorch graphs +to use `cuda_graph_capture_error_mode="thread_local"` (the default). +Unlike PyTorch's strict `"global"` mode, this allows DexSim's Vulkan render +thread to continue making CUDA calls without invalidating capture on the +planner thread. + +If coordinator acquisition times out before recording begins, +`cuda_graph_fallback=True` waits for the active capture to finish and builds a +non-graph backend. An exception after graph recording starts is deliberately +not downgraded: CUDA may have invalidated the process context, so the planner +raises an error and requires a simulator-process restart. The `"global"` and +`"relaxed"` modes remain available for diagnosis, but `"thread_local"` is the +supported renderer-compatible setting. + +cuRobo cannot reset a captured trajectory-optimizer graph when switching +between Cartesian pose goals and joint-space goals. EmbodiChain therefore +caches those two goal types separately and initializes each lazily. Applications +that use only one move type retain one planner backend; using both incurs a +second one-time warmup and its graph-resident memory, but still no subprocess or +second CUDA context. + +The collision world is always auto-generated from live `RigidObject` meshes via +`CuroboWorldCfg.rigid_objects`: the adapter reads each object's mesh +(`get_vertices` / `get_triangles`) and world pose (`get_local_pose`) and writes a +cached cuRobo scene YAML on the first plan, using +`CuroboWorldCfg.obstacle_representation` (`"sphere"` by default for fast +collision queries; use `"cuboid"` for a local-frame AABB placed as an OBB via +the object pose, or `"mesh"` for the exact triangle mesh). +Generated poses are authored in the cuRobo base/world frame, so this is exact +when the robot base sits at the simulator world origin. For obstacles that move +or live in an offset base frame, also declare their names in +`CuroboWorldCfg.dynamic_obstacle_names` and update poses at plan time through +`CuroboPlanOptions.dynamic_obstacle_poses` (provision +`CuroboWorldCfg.collision_cache` before planning). Dynamic updates require the +`"cuboid"` or `"mesh"` representation because sphere fitting expands one object +into multiple independently named obstacles. 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 the generated world +YAML is cloned into one V2 scene per batch row. An empty world (`rigid_objects` +left `None`) 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. + +(curobo-auto-generated-robot-yaml)= +## Auto-generated robot YAML + +On the first plan, the adapter auto-derives the cuRobo robot profile from the +robot's URDF and solver, so nothing robot-specific needs to be hardcoded: + +- `robot_config_path` is produced by `generate_curobo_robot_yaml`, which fits + collision spheres to each link mesh and writes a cuRobo V2 robot YAML. +- The TCP, tool frame, and base link are read from the robot's solver + (`robot._solvers[control_part]`): `tool_frame_name` <- `solver.end_link_name`, + `tool_frame_to_tcp` <- `solver.tcp_xpos`, `base_link_name` <- + `solver.root_link_name`. +- `sim_to_curobo_joint_names` is the identity mapping, since the generated YAML + reuses the simulator's own URDF joint names. + +The generated YAML is cached on disk (default `$XDG_CACHE_HOME/embodichain_curobo` +or `~/.cache/embodichain_curobo`) keyed by the URDF path, URDF content, control +part, tool frame, and fit parameters, so editing the URDF or changing the fit +settings regenerates automatically and subsequent inits reuse the cache. Tune the +fit with `CuroboPlannerCfg.auto_gen` (`fit_type="voxel"` by default for fast +first-generation; `"morphit"` for best quality; `force=True` to bypass the cache). +The default `sphere_density=0.1` keeps the per-link sphere count low (~80 for a +Panda) so planning stays fast; raise it for tighter collision coverage. + +## 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. By default the returned collision-checked samples are +arc-length resampled to the action's `sample_interval` waypoint count (so +`MoveEndEffectorCfg.sample_interval` controls the trajectory length, as for the +other planners); set `CuroboPlannerCfg.preserve_plan_samples=True` to keep +cuRobo's own samples (whose count is derived from `interpolation_dt` and the +trajectory duration). + +~~~python +import torch + +from embodichain.lab.sim.planners import ( + CuroboPlanOptions, + MotionGenOptions, + PlanState, +) + +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. +- Collision worlds are generated from `RigidObject` meshes (cuboid/mesh/sphere) + plus named dynamic pose updates. Arbitrary geometry insertion and removal at + runtime are unsupported. +- The generated collision world assumes a fixed-base robot at the simulator + origin. With a moving base, publish each relevant world obstacle as a named + dynamic pose for every plan; automatic reprojection of static 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 of cuRobo itself and cuRobo V1 compatibility are unsupported. + CPU physics is supported because tensors are transferred to CUDA only for + planning and the resulting trajectory is copied back to the simulation + device. + +## 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 + +# CPU physics with CUDA planning +python examples/sim/planners/curobo_planner.py --headless --sim-device cpu +~~~ + +The demo exports the DexSim `demo_block` into the cuRobo collision world via +`CuroboWorldCfg.rigid_objects` (the robot and world YAMLs are both +auto-generated), prints the result status and trajectory shape, then replays the +returned full-DoF trajectory. CUDA graph capture is enabled by default with the +renderer-compatible `"thread_local"` mode; pass `--no-cuda-graph` to disable 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 d1320be1e..c4b3bcf15 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 cuRobo V2 backend that plans on CUDA and supports either CPU or CUDA physics simulation 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 a CUDA-matched EmbodiChain cuRobo extra 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 4d8b53dba..ed0b7939c 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 CuroboPlanner backend, which plans on CUDA with either CPU or CUDA physics simulation. * **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 (by default) is resampled to the action's `sample_interval`; set `CuroboPlannerCfg.preserve_plan_samples=True` to keep its raw 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/source/quick_start/install.md b/docs/source/quick_start/install.md index 58dc14f36..939064b23 100644 --- a/docs/source/quick_start/install.md +++ b/docs/source/quick_start/install.md @@ -151,6 +151,46 @@ pip install embodichain \ This pulls in `dexsim_engine` (Python package `dexsim`) and the rest of the core dependencies declared in `pyproject.toml`. +## Optional: cuRobo V2 motion planning + +Install a cuRobo extra to use EmbodiChain's CUDA-accelerated, collision-aware +motion planner. cuRobo is intentionally not part of the core dependency set: +select exactly one extra that matches the CUDA version reported by +`nvidia-smi`. + +The normal EmbodiChain environment already provides PyTorch, so prefer one of +the non-`torch` extras: + +| CUDA | Published package | Source checkout | +|------|-------------------|-----------------| +| 12.x | `uv pip install "embodichain[curobo-cu12]" ${PIP_EXTRA_ARGS}` | `uv pip install -e ".[curobo-cu12]" ${PIP_EXTRA_ARGS}` | +| 13.x | `uv pip install "embodichain[curobo-cu13]" ${PIP_EXTRA_ARGS}` | `uv pip install -e ".[curobo-cu13]" ${PIP_EXTRA_ARGS}` | + +For a fresh environment that also needs cuRobo to select and install PyTorch, +use `curobo-cu12-torch` or `curobo-cu13-torch` instead. The same extras work +with `pip`; replace `uv pip install` with `pip install`. + +**Recommended for the current CUDA 12.x EmbodiChain stack:** + +```bash +uv pip install -e ".[curobo-cu12]" \ + --extra-index-url http://pyp.open3dv.site:2345/simple/ \ + --trusted-host pyp.open3dv.site + +python -c "import curobo; print(curobo.__version__)" +pytest --pyargs curobo.tests +``` + +The dependency is installed from NVIDIA's source repository and pinned to the +cuRobo V2 `v0.8.0` release. cuRobo has stricter requirements than the core +EmbodiChain installation: Linux, Python 3.10--3.13, a supported NVIDIA GPU with +at least 4 GB VRAM, and a driver that supports CUDA 12 or newer. See +[NVIDIA's official installation guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html) +for the current compatibility requirements, and see +[cuRobo V2 Planner](../overview/sim/planners/curobo_planner.md) for EmbodiChain +configuration and usage. cuRobo planning always runs on CUDA, but the +SimulationManager physics device may be either CPU or CUDA. + ## Optional: generative simulation (`gensim`) Install the `gensim` extra for SimReady asset pipelines, Blender-based mesh processing, and `pyrender`. The `bpy` wheel is hosted on Blender's index and must be included in the install command. @@ -219,6 +259,7 @@ Press `Ctrl+C` to stop; the script cleans up the simulation on exit. | Docker Vulkan / EGL warnings from `docker_run.sh` | Install host NVIDIA drivers and Vulkan user-space packages; paths must be files under `/etc` or `/usr/share`, not directories. | | Viewer does not open | Export `DISPLAY`, allow X11 access (`xhost +local:` on the host), and ensure `~/.Xauthority` is mounted (the run script does this by default). | | PyTorch / CUDA errors at runtime | Reinstall a PyTorch build that matches your driver/CUDA from [pytorch.org](https://pytorch.org/get-started/locally/). | +| `No module named 'curobo'` | Install exactly one CUDA-matched cuRobo extra, such as `uv pip install -e ".[curobo-cu12]"`, from the EmbodiChain repository root. | | `bpy` install fails | Include the Blender index (`https://download.blender.org/pypi/`) and use Python 3.10 or 3.11. | ## Next steps diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 4e80a4fa8..b11916690 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 d62d50595..212487c44 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 f8c2f88d7..91eff826c 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 bf2b29681..845d6bb5a 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 060c832ea..3f88d35f6 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -212,22 +212,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) diff --git a/embodichain/lab/sim/atomic_actions/primitives/place.py b/embodichain/lab/sim/atomic_actions/primitives/place.py index 9d2cd772f..5198d8730 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/place.py +++ b/embodichain/lab/sim/atomic_actions/primitives/place.py @@ -179,20 +179,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 c30edbba6..82d6fd440 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 e2fc7938e..d48687207 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 @@ -380,38 +388,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 + result.success.to(self.device) 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, ) - # Failed envs hold start qpos + if positions.device != self.device or not torch.isfinite(positions).all(): + logger.log_error( + "MotionGenerator returned non-finite or wrong-device positions", + ValueError, + ) + 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 +501,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 +514,75 @@ 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() + if planner_type == "curobo": + from embodichain.lab.sim.planners.curobo.curobo_planner import ( + CuroboPlanOptions, + ) - return NeuralPlanOptions() + 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 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))``. + """ + 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 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 def plan_joint_traj( self, diff --git a/embodichain/lab/sim/planners/__init__.py b/embodichain/lab/sim/planners/__init__.py index bff37bf78..58cfd5009 100644 --- a/embodichain/lab/sim/planners/__init__.py +++ b/embodichain/lab/sim/planners/__init__.py @@ -18,4 +18,6 @@ from .base_planner import * from .toppra_planner import * from .neural_planner import * +from .curobo.curobo_yaml import * +from .curobo.curobo_planner import * from .motion_generator import * diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index c71866f34..45612bd39 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -155,6 +155,59 @@ 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. + """ + + 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() + + 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/curobo/curobo_planner.py b/embodichain/lab/sim/planners/curobo/curobo_planner.py new file mode 100644 index 000000000..d4558931b --- /dev/null +++ b/embodichain/lab/sim/planners/curobo/curobo_planner.py @@ -0,0 +1,2219 @@ +# ---------------------------------------------------------------------------- +# 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 hashlib +import importlib +import os +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from copy import deepcopy +from dataclasses import 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 pose_inv, quat_from_matrix + +from embodichain.lab.sim.planners.base_planner import ( + BasePlanner, + BasePlannerCfg, + PlanOptions, + validate_plan_options, +) +from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState + +if TYPE_CHECKING: + from typing import Any + + from embodichain.lab.sim.objects import RigidObject + +__all__ = [ + "CuroboAutoGenCfg", + "CuroboPlanOptions", + "CuroboPlanner", + "CuroboPlannerCfg", + "CuroboWorldCfg", +] + + +# cuRobo V2 installation extras documented at NVIDIA's installation page. +_CUROBO_INSTALL_URL = ( + "https://nvlabs.github.io/curobo/latest/getting-started/installation.html" +) + +# Bumped whenever the auto-generated robot-YAML schema/logic changes so that +# cached YAMLs from an older generator are regenerated instead of reused. v2: +# exclude URDF mimic joints from cspace/lock_joints (cuRobo folds them into +# their active joint and raises KeyError when locking one). +_CUROBO_ROBOT_YAML_GENERATOR_VERSION = "v2" + +# cuRobo 0.8 does not expose PyTorch's CUDA stream-capture error mode. The +# temporary adapter below therefore replaces ``torch.cuda.graph`` only while +# cuRobo can lazily record graphs. Serialize that small process-wide patch. +_TORCH_CUDA_GRAPH_PATCH_LOCK = threading.Lock() + + +@dataclass +class _CuroboProfile: + """Auto-derived cuRobo robot profile for one control part (internal). + + Produced by :meth:`CuroboPlanner._materialize_profile` from the robot's URDF + and IK solver - never user-configured. The cuRobo robot YAML is always + auto-generated from the URDF (see :class:`CuroboAutoGenCfg`), so the simulator + and cuRobo share the same joint names and the joint mapping is identity. + """ + + robot_config_path: str + """Cached path of the auto-generated cuRobo robot YAML.""" + + sim_to_curobo_joint_names: dict[str, str] + """Simulator -> cuRobo joint-name mapping (identity for the auto-gen YAML).""" + + tool_frame_name: str | None = None + """cuRobo tool frame (a URDF link 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. + + ``None`` means the tool frame is already the TCP (the common auto-derived + case, where the solver's ``end_link_name`` is the TCP). + """ + + base_link_name: str | None = None + """cuRobo robot base link, validated against the loaded V2 model.""" + + sim_base_link_name: str | None = None + """Simulator link physically equivalent to the control-part base.""" + + sim_base_to_curobo_base: list[list[float]] | None = None + """Fixed transform from the simulator base to the cuRobo base (``None``=coincide).""" + + +class _RigidObjectRefList(list): + """A list of live ``RigidObject`` handles that survives ``@configclass`` deepcopy. + + ``@configclass`` deepcopies every field on construction, but live dexsim + objects hold non-pickleable C++ handles (e.g. ``dexsim.World``). This + ``list`` subclass overrides ``__deepcopy__`` to share the object references + instead of cloning them, so ``CuroboWorldCfg(rigid_objects=[...])`` works. + """ + + def __deepcopy__(self, memo: dict) -> "_RigidObjectRefList": # noqa: ARG002 + return _RigidObjectRefList(self) + + +@configclass +class CuroboWorldCfg: + """Static collision-world configuration for the cuRobo backend. + + The collision world is always auto-generated from live :class:`RigidObject` + meshes (see :attr:`rigid_objects`); there is no external scene-YAML path. + """ + + rigid_objects: list[RigidObject] | None = None + """Live :class:`RigidObject` obstacles to bake into the auto-generated world YAML. + + The adapter reads each object's mesh (``get_vertices`` / ``get_triangles``) + and world pose (``get_local_pose``) and writes a cuRobo V2 scene YAML (cached + on disk by content hash). Poses are written in the cuRobo world/base frame, + so this is exact when the robot base sits at the simulator world origin. For + obstacles that move or live in an offset base frame, also list their names in + :attr:`dynamic_obstacle_names` to update poses at plan time. ``None`` yields an + initially empty collision world. + """ + + obstacle_representation: str = "sphere" + """Collision representation used when generating the YAML from :attr:`rigid_objects`. + + ``"sphere"`` (default) fits spheres with cuRobo's + ``fit_spheres_to_mesh`` (fast collision queries, approximate, and requires + CUDA + cuRobo + trimesh). ``"cuboid"`` emits a local-frame AABB per object, + placed as an OBB via the object pose. ``"mesh"`` emits the full triangle + mesh (exact, no CUDA). + """ + + 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 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. The generated world YAML is cloned for every row. + """ + + def __post_init__(self) -> None: + # Wrap live RigidObjects so the @configclass field-deepcopy (run right + # after this by custom_post_init) shares references instead of trying to + # pickle non-pickleable C++ dexsim handles held by each RigidObject. + if self.rigid_objects is not None and not isinstance( + self.rigid_objects, _RigidObjectRefList + ): + self.rigid_objects = _RigidObjectRefList(self.rigid_objects) + + +@configclass +class CuroboAutoGenCfg: + """Auto-generation of the cuRobo robot YAML from the robot's URDF. + + The adapter generates a cuRobo robot configuration YAML from the robot's URDF + (fitting collision spheres to each link mesh) on the first plan and caches it + on disk so subsequent inits skip regeneration. The TCP, tool frame, and base + link are read from the robot's solver, so nothing robot-specific needs to be + hardcoded. + """ + + cache_dir: str | None = None + """Directory for cached robot YAMLs. + + ``None`` (default) uses ``$XDG_CACHE_HOME/embodichain_curobo`` or + ``~/.cache/embodichain_curobo``. The cache key hashes the generator version, + URDF path, URDF content, control part, tool frame, and fit parameters, so + editing the URDF, changing the fit settings, or a generator update + regenerates automatically. + """ + + fit_type: str = "voxel" + """cuRobo sphere-fit strategy for auto-generation: ``"voxel"`` (default, + fast), ``"morphit"`` (best, slower), or ``"surface"`` (crude).""" + + num_spheres: int | None = None + """Per-link sphere count. ``None`` auto-estimates from bounding-box volume + scaled by :attr:`sphere_density`.""" + + sphere_density: float = 0.1 + """Multiplier on the auto-estimated per-link sphere count (ignored when + :attr:`num_spheres` is set). + + The cuRobo volume-based estimate over-fits at ``1.0`` (~668 spheres for a + Franka Panda, making planning pathologically slow). ``0.1`` (default) yields + ~50-100 spheres - enough coverage for collision-aware planning while keeping + each plan fast. Increase for tighter coverage on complex robots. + """ + + surface_radius: float = 0.005 + """Fixed radius used only by the ``surface`` strategy.""" + + iterations: int = 200 + """Adam iterations for the ``morphit`` strategy.""" + + collision_sphere_buffer: float = 0.0 + """Padding added to every fitted sphere's radius (m).""" + + force: bool = False + """Bypass the cache and regenerate the robot YAML on the next plan.""" + + +@configclass +class CuroboPlannerCfg(BasePlannerCfg): + """Configuration for the cuRobo V2 planner backend. + + cuRobo runs in the simulator process so it reuses the existing CUDA context + instead of keeping a spawned Python process and a second CUDA context alive. + CUDA graphs are enabled by default with renderer-compatible thread-local + capture. Both the cuRobo robot YAML and the collision-world YAML are + auto-generated internally (from the robot's URDF and from + :attr:`world.rigid_objects` respectively); no external YAML is used. The + per-control-part profile is auto-derived from the robot's solver at plan time. + """ + + planner_type: str = "curobo" + + world: CuroboWorldCfg = CuroboWorldCfg() + """Collision-world configuration (auto-generated from ``RigidObject`` meshes).""" + + auto_gen: CuroboAutoGenCfg = CuroboAutoGenCfg() + """Auto-generation settings for the cuRobo robot YAML from the robot's URDF.""" + + 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 together with the live simulator base pose to convert + simulator-world Cartesian goals and dynamic obstacle poses into cuRobo's base + frame. ``None`` (default) means the two base frames coincide - the common + case, since the auto-generated robot YAML is rooted at the URDF base link the + solver reports. Set this only when the simulator base and the URDF base use + different fixed frame conventions. + """ + + 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.""" + + cuda_device: str | int | torch.device | None = None + """CUDA device used exclusively by cuRobo. + + ``None`` uses the simulator GPU when physics runs on CUDA, otherwise the + current PyTorch CUDA device. An integer selects ``cuda:``. CPU + physics is supported, but cuRobo itself always runs on CUDA. + """ + + use_cuda_graph: bool = True + """Whether cuRobo may capture CUDA graphs in the simulator process. + + ``True`` enables the renderer-compatible fast path by default. Capture is + serialized with DexSim Newton captures on the same CUDA device and fenced + with device synchronizations. cuRobo's PyTorch graph captures use + :attr:`cuda_graph_capture_error_mode`, whose ``"thread_local"`` default + permits the DexSim render thread to continue submitting CUDA work. Set this + to ``False`` to reduce one-time initialization and graph-resident memory. + """ + + cuda_graph_fallback: bool = True + """Use non-graph mode if the capture coordinator times out before capture. + + An error after capture starts is never downgraded in-process because CUDA + may leave the context in an invalid state. Such errors are raised and the + simulator process must be restarted. + """ + + cuda_graph_capture_error_mode: str = "thread_local" + """PyTorch CUDA stream-capture error mode used by cuRobo. + + ``"thread_local"`` isolates capture from CUDA calls made by DexSim's render + thread and is the recommended mode for the in-process planner. ``"global"`` + retains PyTorch's strict default and is expected to conflict with an active + renderer. ``"relaxed"`` disables additional capture-safety checks and + should only be used for diagnosis. + """ + + capture_acquire_timeout: float | None = 2.0 + """Seconds to wait for the per-device capture coordinator. + + ``None`` waits indefinitely. After a finite timeout, graph mode falls back + to non-graph mode, then waits for the active capture to finish before + launching planner initialization kernels. + """ + + capture_wait_log_interval: float | None = 10.0 + """Seconds between coordinator wait logs; ``None`` disables them.""" + + interpolation_dt: float = 0.025 + """Interpolation step (seconds) used by cuRobo and as a dt fallback.""" + + preserve_plan_samples: bool = False + """Whether callers must retain cuRobo's raw collision-checked samples exactly. + + When ``False`` (default), :class:`~embodichain.lab.sim.atomic_actions.trajectory.TrajectoryBuilder` + resamples the returned trajectory to the atomic action's ``sample_interval`` + waypoint count - matching the documented contract of + :class:`~embodichain.lab.sim.atomic_actions.primitives.move_end_effector.MoveEndEffectorCfg.sample_interval` + and the other planners. The resample is arc-length piecewise-linear along + cuRobo's joint-space path, so the collision-free path is preserved; only the + sample density changes (cuRobo's own count is derived from + :attr:`interpolation_dt` and the trajectory duration, e.g. ~82 for a 2 s + plan at 0.025 s). + + When ``True``, the builder returns cuRobo's own samples unchanged. Use this + when you need cuRobo's exact time-parameterized, collision-checked samples + rather than a fixed waypoint count. + """ + + warmup_iterations: int = 1 + """cuRobo warmup iterations run once per cached in-process planner. + + Set to ``0`` to skip warmup when CUDA graphs are disabled. Graph mode always + runs at least one coordinated warmup because otherwise the first real plan + would perform an uncoordinated lazy capture. + """ + + +@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 _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) + # 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 + + +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 _ensure_warp_torch_compat() -> None: + """Restore ``warp.torch`` for cuRobo 0.8 when using Warp 1.13 or newer.""" + import sys + + import warp as wp + + if hasattr(wp, "torch"): + return + try: + import warp._src.torch as torch_interop + except ImportError: + return + wp.torch = torch_interop # type: ignore[attr-defined] + sys.modules["warp.torch"] = torch_interop + + +class _LocalCaptureCoordinator: + """Fallback per-device capture lock when DexSim's coordinator is unavailable.""" + + _instance: "_LocalCaptureCoordinator | None" = None + _instance_lock = threading.Lock() + + def __init__(self) -> None: + self._locks: dict[str, threading.Lock] = {} + self._locks_guard = threading.Lock() + + @classmethod + def get(cls) -> "_LocalCaptureCoordinator": + """Return the process-wide fallback coordinator.""" + with cls._instance_lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def acquire_for_capture( + self, + device: str, + owner: object, # noqa: ARG002 - API-compatible with DexSim coordinator + timeout: float | None = None, + wait_log_interval: float | None = None, # noqa: ARG002 + ) -> bool: + """Acquire the lock for ``device`` within ``timeout``.""" + with self._locks_guard: + lock = self._locks.setdefault(str(device), threading.Lock()) + if timeout is None: + lock.acquire() + return True + return lock.acquire(timeout=max(0.0, float(timeout))) + + def release_for_capture( + self, + device: str, + owner: object, # noqa: ARG002 - API-compatible with DexSim coordinator + ) -> None: + """Release the lock for ``device`` when held.""" + with self._locks_guard: + lock = self._locks.get(str(device)) + if lock is not None and lock.locked(): + lock.release() + + +class _CaptureOwner: + """Weak-referenceable identity for one backend capture attempt.""" + + +@contextmanager +def _torch_cuda_graph_capture_mode(mode: str) -> Iterator[None]: + """Temporarily force a capture mode for cuRobo's PyTorch graph contexts. + + cuRobo 0.8 creates every graph through ``torch.cuda.graph`` but does not + forward its ``capture_error_mode`` argument. The patch is deliberately + scoped to planner warmup/planning and restored even when capture raises. + """ + valid_modes = ("global", "thread_local", "relaxed") + if mode not in valid_modes: + raise ValueError( + "CuroboPlannerCfg.cuda_graph_capture_error_mode must be one of " + f"{valid_modes}, got {mode!r}." + ) + + with _TORCH_CUDA_GRAPH_PATCH_LOCK: + original_graph = torch.cuda.graph + + def graph_with_capture_mode(*args: "Any", **kwargs: "Any") -> "Any": + kwargs["capture_error_mode"] = mode + return original_graph(*args, **kwargs) + + torch.cuda.graph = graph_with_capture_mode # type: ignore[assignment] + try: + yield + finally: + torch.cuda.graph = original_graph # type: ignore[assignment] + + +def _get_capture_coordinator() -> "Any": + """Return DexSim Newton's shared capture coordinator when available.""" + try: + coordinator_mod = importlib.import_module( + "dexsim.engine.newton_physics.capture_coordinator" + ) + except (ImportError, AttributeError): + logger.log_warning( + "DexSim CaptureCoordinator is unavailable; cuRobo CUDA graph " + "capture will use a local per-device lock." + ) + return _LocalCaptureCoordinator.get() + return coordinator_mod.CaptureCoordinator.get() + + +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. + """ + # cuRobo 0.8 references ``wp.torch.*``, which Warp >= 1.13 relocated. + _ensure_warp_torch_compat() + try: + 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. " + "From the EmbodiChain repository root, install the CUDA-matched " + "extra, e.g. `pip install -e '.[curobo-cu12]'` for CUDA 12.x or " + "`pip install -e '.[curobo-cu13]'` for CUDA 13.x " + "(also `.[curobo-cu12-torch]` / `.[curobo-cu13-torch]`). " + f"See {_CUROBO_INSTALL_URL} for details." + ) from exc + return SimpleNamespace( + MotionPlanner=planner_mod.MotionPlanner, + MotionPlannerCfg=planner_mod.MotionPlannerCfg, + BatchMotionPlanner=batch_mod.BatchMotionPlanner, + JointState=types_mod.JointState, + Pose=types_mod.Pose, + GoalToolPose=types_mod.GoalToolPose, + DeviceCfg=types_mod.DeviceCfg, + ) + + +def _resolve_curobo_device( + configured_device: str | int | torch.device | None, + simulation_device: torch.device, +) -> torch.device: + """Resolve cuRobo's concrete CUDA device independently of physics.""" + if not torch.cuda.is_available(): + raise RuntimeError( + "cuRobo V2 requires CUDA even when SimulationManager uses CPU " + "physics, but torch.cuda.is_available() is False." + ) + + if configured_device is None: + if simulation_device.type == "cuda" and simulation_device.index is not None: + device = simulation_device + else: + device = torch.device("cuda", torch.cuda.current_device()) + elif isinstance(configured_device, int): + device = torch.device("cuda", configured_device) + else: + device = torch.device(configured_device) + if device.type != "cuda": + raise ValueError( + "CuroboPlannerCfg.cuda_device must select a CUDA device, " + f"got {configured_device!r}." + ) + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + + assert device.index is not None + device_count = torch.cuda.device_count() + if device.index < 0 or device.index >= device_count: + raise RuntimeError( + f"cuRobo CUDA device index {device.index} is unavailable; " + f"torch reports {device_count} CUDA device(s)." + ) + return device + + +# ============================================================================= +# CuroboPlanner +# ============================================================================= + + +class CuroboPlanner(BasePlanner): + r"""cuRobo V2 collision-aware motion-planning backend. + + The planner lazily imports cuRobo V2 at construction time and builds a + ``MotionPlanner`` (single environment) or ``BatchMotionPlanner`` (batched + environments) in the simulator process. Backends are cached per control + part, batch size, collision-world mode, and goal type, so there is no helper + process, IPC tensor copy, or second CUDA context. + + CUDA graphs are enabled by default. Initialization uses DexSim Newton's + per-device capture coordinator, synchronizes the CUDA device before and + after warmup, and records PyTorch graphs in + ``"thread_local"`` capture mode so the DexSim renderer can continue issuing + CUDA calls from its own thread. A coordinator timeout may safely downgrade + before capture begins; a failure during capture is raised because the CUDA + context may already be invalid. + + 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``). + By default the returned collision-checked samples are arc-length resampled to + the action's ``sample_interval`` waypoint count + (``preserve_plan_samples=False``); set + :attr:`CuroboPlannerCfg.preserve_plan_samples=True` to keep cuRobo's own + samples unchanged. + + Args: + cfg: Configuration for the cuRobo planner. + + Raises: + ImportError: If cuRobo V2 is not installed. + RuntimeError: If CUDA is unavailable for cuRobo. + ValueError: If ``robot_uid`` is missing or the robot is not found. + """ + + preinterpolate_targets = False + supports_joint_move = True + + @property + def preserve_plan_samples(self) -> bool: + """Whether callers must retain this planner's raw samples exactly. + + Mirrors :attr:`CuroboPlannerCfg.preserve_plan_samples`; read by + :class:`~embodichain.lab.sim.atomic_actions.trajectory.TrajectoryBuilder` + to decide whether to resample the returned trajectory to the action's + ``sample_interval``. + """ + return self.cfg.preserve_plan_samples + + def __init__(self, cfg: CuroboPlannerCfg) -> None: + super().__init__(cfg) + self.cfg: CuroboPlannerCfg = cfg + self.device = torch.device(self.robot.device) + self._curobo_device = _resolve_curobo_device( + cfg.cuda_device, + self.device, + ) + # cuRobo and Warp contain a few current-device-sensitive initialization + # paths, so select the dedicated planning GPU before importing them. + torch.cuda.set_device(self._curobo_device) + self._bindings = _require_curobo() + self._backend_cache: dict[tuple[str, int, bool, MoveType], "_CuroboBackend"] = ( + {} + ) + world_cfg = cfg.world + if world_cfg.obstacle_representation not in ("cuboid", "mesh", "sphere"): + logger.log_error( + "CuroboWorldCfg.obstacle_representation must be 'cuboid', 'mesh', " + f"or 'sphere', got {world_cfg.obstacle_representation!r}.", + ValueError, + ) + if ( + world_cfg.dynamic_obstacle_names + and world_cfg.obstacle_representation == "sphere" + ): + logger.log_error( + "Dynamic obstacle updates require the 'cuboid' or 'mesh' world " + "representation. Sphere fitting expands one RigidObject into " + "multiple independent obstacles that cannot be updated by the " + "original object name.", + ValueError, + ) + if cfg.warmup_iterations < 0: + logger.log_error( + "CuroboPlannerCfg.warmup_iterations must be non-negative.", + ValueError, + ) + if ( + cfg.capture_acquire_timeout is not None + and cfg.capture_acquire_timeout < 0.0 + ): + logger.log_error( + "CuroboPlannerCfg.capture_acquire_timeout must be non-negative " + "or None.", + ValueError, + ) + if cfg.cuda_graph_capture_error_mode not in ( + "global", + "thread_local", + "relaxed", + ): + logger.log_error( + "CuroboPlannerCfg.cuda_graph_capture_error_mode must be " + "'global', 'thread_local', or 'relaxed'; got " + f"{cfg.cuda_graph_capture_error_mode!r}.", + ValueError, + ) + + 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``. + + ``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`` + 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. ``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``. + """ + if not target_states: + return PlanResult( + success=torch.zeros(0, dtype=torch.bool, device=self.device), + positions=None, + ) + control_part = self._resolve_control_part(options) + start = self._resolve_start_qpos(options.start_qpos, control_part) + move_types = {target.move_type for target in target_states} + unsupported = move_types.difference((MoveType.EEF_MOVE, MoveType.JOINT_MOVE)) + if unsupported: + logger.log_error( + f"cuRobo does not support move types {sorted(str(x) for x in unsupported)}.", + ValueError, + ) + backends = { + move_type: self._get_backend( + control_part, + start.shape[0], + move_type, + ) + for move_type in move_types + } + transform_backend = backends.get( + MoveType.EEF_MOVE, next(iter(backends.values())) + ) + # Compute the live sim base pose + its inverse once per plan and reuse it + # across every EEF segment and every dynamic-obstacle update (the robot + # does not move during planning), instead of re-querying get_link_pose + + # re-inverting per segment / obstacle. Skipped for pure + # joint-move plans with no dynamic obstacles, which never need it. + needs_base_pose = bool(options.dynamic_obstacle_poses) or any( + t.move_type == MoveType.EEF_MOVE for t in target_states + ) + sim_base_pose_inv = ( + pose_inv(self._get_sim_base_pose(transform_backend, start.shape[0])) + if needs_base_pose + else None + ) + for backend in backends.values(): + self.update_dynamic_obstacles( + options.dynamic_obstacle_poses, backend, sim_base_pose_inv + ) + return self._plan_segments( + target_states, start, backends, options, sim_base_pose_inv + ) + + # ------------------------------------------------------------------ + # Profile / start resolution + # ------------------------------------------------------------------ + + def _resolve_control_part(self, options: CuroboPlanOptions) -> str: + """Resolve and validate the requested control part against the robot.""" + control_part = options.control_part + if control_part is None: + logger.log_error("CuroboPlanOptions.control_part is required.", ValueError) + control_parts = getattr(self.robot, "control_parts", None) or {} + if control_part not in control_parts: + logger.log_error( + f"Robot '{self.cfg.robot_uid}' has no control part '{control_part}'. " + f"Available control parts: {sorted(control_parts)}.", + ValueError, + ) + return 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._curobo_device + ) + if start_qpos.dim() == 1: + start_qpos = start_qpos.unsqueeze(0) + return start_qpos + + # ------------------------------------------------------------------ + # Backend construction / caching + # ------------------------------------------------------------------ + + def _materialize_multi_env_scene_model( + self, world_config_path: str | None, batch_size: int + ) -> list[dict]: + """Return one independent cuRobo scene mapping for every batch row.""" + 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: + 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, + control_part: str, + batch_size: int, + planning_mode: MoveType = MoveType.EEF_MOVE, + ) -> "_CuroboBackend": + """Return a cached in-process backend for one goal-buffer shape.""" + multi_env = bool(self.cfg.world.multi_env) + key = (control_part, int(batch_size), multi_env, planning_mode) + if key in self._backend_cache: + return self._backend_cache[key] + + profile = self._materialize_profile(control_part) + 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 + ) + world_config_path = ( + self._auto_generate_world_yaml(world_cfg) + if world_cfg.rigid_objects + else None + ) + scene_model: str | list[dict] | None = world_config_path + if multi_env: + scene_model = self._materialize_multi_env_scene_model( + world_config_path, int(batch_size) + ) + + use_cuda_graph = bool(self.cfg.use_cuda_graph) + coordinator = None + capture_owner = _CaptureOwner() + capture_acquired = False + if use_cuda_graph: + coordinator = _get_capture_coordinator() + capture_acquired = coordinator.acquire_for_capture( + str(self._curobo_device), + capture_owner, + timeout=self.cfg.capture_acquire_timeout, + wait_log_interval=self.cfg.capture_wait_log_interval, + ) + if not capture_acquired: + if not self.cfg.cuda_graph_fallback: + raise RuntimeError( + "Timed out waiting for coordinated cuRobo CUDA graph " + f"capture on {self._curobo_device}." + ) + logger.log_warning( + "Timed out waiting for coordinated cuRobo CUDA graph capture " + f"on {self._curobo_device}; waiting for the active capture " + "to finish, then using non-graph mode." + ) + use_cuda_graph = False + # Non-graph initialization still launches CUDA kernels. Wait + # for the active capture to finish so those launches cannot + # invalidate a peer's stream capture. + capture_acquired = coordinator.acquire_for_capture( + str(self._curobo_device), + capture_owner, + timeout=None, + wait_log_interval=self.cfg.capture_wait_log_interval, + ) + if not capture_acquired: + raise RuntimeError( + "Unable to coordinate safe non-graph cuRobo " + f"initialization on {self._curobo_device}." + ) + + backend: _CuroboBackend + try: + backend = self._build_backend( + control_part=control_part, + batch_size=int(batch_size), + profile=profile, + sim_joint_names=sim_joint_names, + scene_model=scene_model, + collision_cache=collision_cache, + use_cuda_graph=use_cuda_graph, + planning_mode=planning_mode, + ) + try: + self._warmup_backend(backend) + except Exception as exc: + self._close_planner(backend.planner) + if not use_cuda_graph: + raise + raise RuntimeError( + "cuRobo CUDA graph warmup failed after capture may have " + "started. The CUDA context may now be invalid, so an " + "in-process non-graph fallback is unsafe; restart the " + "simulator process. Original error: " + f"{exc}" + ) from exc + finally: + if capture_acquired and coordinator is not None: + try: + coordinator.release_for_capture( + str(self._curobo_device), capture_owner + ) + except Exception as exc: + logger.log_warning( + f"cuRobo capture coordinator release failed: {exc}" + ) + + self._backend_cache[key] = backend + logger.log_info( + f"cuRobo in-process backend ready for '{control_part}' " + f"(batch={batch_size}, mode={planning_mode.name}, " + f"cuda_graph={backend.use_cuda_graph})." + ) + return backend + + def _build_backend( + self, + *, + control_part: str, + batch_size: int, + profile: _CuroboProfile, + sim_joint_names: list[str], + scene_model: str | list[dict] | None, + collision_cache: dict[str, int | dict[str, int | float | list[float]]] | None, + use_cuda_graph: bool, + planning_mode: MoveType, + ) -> "_CuroboBackend": + """Construct and validate one cuRobo planner on the selected CUDA device.""" + with torch.cuda.device(self._curobo_device): + planner_cfg = self._bindings.MotionPlannerCfg.create( + robot=profile.robot_config_path, + scene_model=scene_model, + collision_cache=collision_cache, + device_cfg=self._bindings.DeviceCfg(device=self._curobo_device), + max_batch_size=batch_size, + multi_env=bool(self.cfg.world.multi_env), + optimizer_collision_activation_distance=( + self.cfg.collision_activation_distance + ), + use_cuda_graph=use_cuda_graph, + ) + # cuRobo 0.8 reads interpolation_dt from the trajectory optimizer + # config rather than accepting it in MotionPlannerCfg.create(). + planner_cfg.trajopt_solver_config.interpolation_dt = float( + self.cfg.interpolation_dt + ) + planner = ( + self._bindings.MotionPlanner(planner_cfg) + if batch_size == 1 + else self._bindings.BatchMotionPlanner(planner_cfg) + ) + + try: + 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) + except Exception: + self._close_planner(planner) + raise + return _CuroboBackend( + planner=planner, + control_part=control_part, + sim_joint_names=sim_joint_names, + tool_frame=tool_frame, + profile=profile, + batch_size=batch_size, + use_cuda_graph=use_cuda_graph, + planning_mode=planning_mode, + ) + + def _warmup_backend(self, backend: "_CuroboBackend") -> None: + """Warm one goal type without forcing cuRobo to reset captured graphs. + + cuRobo uses structurally different trajectory-optimizer goal buffers for + pose and c-space solves. Its CUDA backend cannot reset captured graphs, + so each cached backend is warmed only for its declared planning mode. + """ + iterations = int(self.cfg.warmup_iterations) + if not backend.use_cuda_graph and iterations == 0: + return + if backend.use_cuda_graph: + iterations = max(iterations, 1) + self._synchronize_device() + capture_mode = ( + _torch_cuda_graph_capture_mode(self.cfg.cuda_graph_capture_error_mode) + if backend.use_cuda_graph + else nullcontext() + ) + with capture_mode, torch.cuda.device(self._curobo_device): + planner = backend.planner + default_position = planner.default_joint_state.position + if default_position.dim() == 1: + default_position = default_position.unsqueeze(0) + if default_position.shape[0] == 1 and backend.batch_size > 1: + default_position = default_position.expand( + backend.batch_size, -1 + ).clone() + current_state = self._bindings.JointState.from_position( + default_position, + joint_names=list(planner.joint_names), + ) + original_exit_early = planner.ik_solver.config.exit_early + planner.ik_solver.config.exit_early = False + try: + for _ in range(iterations): + goal_state = current_state.clone() + goal_state.position[..., 0] += 0.2 + if backend.planning_mode == MoveType.EEF_MOVE: + goal = planner.compute_kinematics( + goal_state + ).tool_poses.as_goal() + planner.plan_pose( + goal, + current_state, + max_attempts=1, + enable_graph_attempt=1, + ) + elif backend.planning_mode == MoveType.JOINT_MOVE: + planner.plan_cspace( + goal_state, + current_state, + max_attempts=1, + enable_graph_attempt=1, + ) + else: # pragma: no cover - validated before backend creation + raise ValueError( + f"Unsupported cuRobo warmup mode {backend.planning_mode}." + ) + planner.reset_seed() + + graph_planner = getattr(planner, "graph_planner", None) + if backend.use_cuda_graph and graph_planner is not None: + graph_planner.warmup(num_warmup_iterations=iterations) + finally: + planner.ik_solver.config.exit_early = original_exit_early + self._synchronize_device() + + def _synchronize_device(self) -> None: + """Synchronize only the CUDA device used by this planner.""" + torch.cuda.synchronize(self._curobo_device) + + @staticmethod + def _close_planner(planner: "Any") -> None: + """Best-effort release of a cuRobo planner's graph resources.""" + close_fn = getattr(planner, "close", None) or getattr(planner, "destroy", None) + if close_fn is not None: + try: + close_fn() + except Exception: + pass + + def _validate_profile_joint_names( + self, + profile: _CuroboProfile, + sim_joint_names: list[str], + curobo_joint_names: list[str], + ) -> None: + """Validate the auto-derived joint 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, + ) + mapped_set = set(mapped_names) + unmapped = [name for name in curobo_joint_names if name not in mapped_set] + 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, + ) + + def _resolve_tool_frame(self, profile: _CuroboProfile, planner: "Any") -> str: + """Resolve and validate the 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 + + @staticmethod + def _validate_base_link_name(profile: _CuroboProfile, planner: "Any") -> None: + """Ensure the auto-derived base link matches the loaded V2 model.""" + expected = profile.base_link_name + if expected is None: + return + actual = getattr(getattr(planner, "kinematics", None), "base_link", None) + if actual is None: + logger.log_error( + "cuRobo planner did not expose kinematics.base_link, so " + f"base_link_name={expected!r} cannot be validated.", + ValueError, + ) + if actual != expected: + logger.log_error( + f"Auto-derived base_link_name={expected!r} does not match the " + f"loaded cuRobo V2 base link {actual!r}.", + ValueError, + ) + + def _materialize_profile(self, control_part: str) -> _CuroboProfile: + """Auto-derive the cuRobo profile for ``control_part`` from the robot. + + Reads the tool frame, TCP offset, and base link from the control part's + IK solver, builds the identity simulator->cuRobo joint mapping (the + auto-generated robot YAML reuses the URDF joint names), and generates + the cuRobo robot YAML from the URDF. Nothing robot-specific is hardcoded. + """ + robot = self.robot + assert ( + robot is not None + ), "cuRobo planner has no robot; cannot materialize the profile." + solver = None + solvers = getattr(robot, "_solvers", None) or {} + if solvers and control_part in solvers: + solver = solvers[control_part] + + # Tool frame: prefer the solver's end link (the TCP), else the control + # part's last link. Auto-generation needs a concrete tool frame. + tool_frame = ( + getattr(solver, "end_link_name", None) if solver is not None else None + ) + if tool_frame is None: + part_links = robot.get_control_part_link_names(control_part) or [] + if not part_links: + logger.log_error( + f"Control part {control_part!r} has no solver end_link_name and " + "no links; cannot derive a cuRobo tool frame.", + ValueError, + ) + tool_frame = part_links[-1] + + # TCP offset: only when the solver's tool frame is not itself the TCP. + tool_frame_to_tcp = None + if solver is not None: + tcp_xpos = getattr(solver, "tcp_xpos", None) + if tcp_xpos is not None: + tool_frame_to_tcp = tcp_xpos.tolist() + + # cuRobo's base is the auto-generated YAML's ``base_link`` (the URDF root + # link), NOT the solver's control-part root. For robots whose control + # part spans the whole arm (franka, ur) the two coincide, but for a + # control part that is a sub-chain of a larger robot - e.g. w1 + # ``right_arm`` whose root ``right_arm_base`` hangs off a locked torso - + # they differ. Cartesian goals and dynamic obstacle poses are converted + # into this base frame (see :meth:`_sim_world_to_curobo_base_pose`), so it + # must match the frame cuRobo actually plans in; otherwise cuRobo receives + # a goal expressed in the control-part base and interprets it in the URDF + # root, planning to a wrong pose. + solver_base_link = ( + getattr(solver, "root_link_name", None) if solver is not None else None + ) + + sim_joints = self._resolve_sim_joint_names(control_part) + sim_to_curobo = {j: j for j in sim_joints} + + robot_config_path = self._auto_generate_robot_yaml(control_part, tool_frame) + base_link = self._read_curobo_base_link(robot_config_path) or solver_base_link + sim_base_link = base_link + + return _CuroboProfile( + robot_config_path=robot_config_path, + sim_to_curobo_joint_names=sim_to_curobo, + tool_frame_name=tool_frame, + tool_frame_to_tcp=tool_frame_to_tcp, + base_link_name=base_link, + sim_base_link_name=sim_base_link, + sim_base_to_curobo_base=self.cfg.sim_base_to_curobo_base, + ) + + @staticmethod + def _read_curobo_base_link(robot_yaml_path: str) -> str | None: + """Return cuRobo's ``base_link`` from an auto-generated robot YAML. + + The YAML's ``robot_cfg.kinematics.base_link`` is the URDF root link cuRobo + roots its kinematics at - the frame Cartesian goals must be expressed in. + Reading it back (rather than assuming it equals the solver's control-part + root) keeps the parent's frame conversion in sync with cuRobo's actual + model for robots whose control part is a sub-chain of a larger URDF. + + Args: + robot_yaml_path: Path to the cached cuRobo robot YAML. + + Returns: + The base link name, or ``None`` if the YAML cannot be read. + """ + try: + import yaml + + with open(robot_yaml_path, "r") as fh: + data = yaml.safe_load(fh) + return data["robot_cfg"]["kinematics"]["base_link"] + except Exception: # noqa: BLE001 - fall back to the solver root upstream + return None + + def _auto_generate_robot_yaml( + self, control_part: str, tool_frame: str | None + ) -> str: + """Return a cached cuRobo robot YAML path, generating it from the URDF if needed.""" + from .curobo_yaml import generate_curobo_robot_yaml + + robot = self.robot + assert ( + robot is not None + ), "cuRobo planner has no robot; cannot auto-generate its YAML." + auto = self.cfg.auto_gen + # cuRobo's robot YAML is generated from the *assembled* URDF + # (robot.cfg.fpath), which includes every mounted component (arm + + # gripper). A solver's ``urdf_path`` may be a sub-chain URDF (e.g. the + # UR arm's bare UR10 URDF hardcoded in URSolverCfg) that omits the + # gripper; keying the cache on it would reuse a stale, gripper-less YAML + # even after the gripper is attached, and would not invalidate when the + # gripper changes. Use robot.cfg.fpath for both the cache key and + # generation so they stay consistent and the gripper links are included. + urdf_path = robot.cfg.fpath + cache_dir = auto.cache_dir or os.path.join( + os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), + "embodichain_curobo", + ) + cache_key = self._robot_yaml_cache_key( + urdf_path, control_part, tool_frame, auto + ) + cache_path = os.path.join(cache_dir, f"{cache_key}.yml") + if not auto.force and os.path.exists(cache_path): + logger.log_info(f"cuRobo robot YAML cache hit: {cache_path}") + return cache_path + logger.log_info( + f"Auto-generating cuRobo robot YAML from URDF ({urdf_path}) -> {cache_path}" + ) + return generate_curobo_robot_yaml( + robot, + control_part, + cache_path, + tool_frame=tool_frame, + urdf_path=urdf_path, + fit_type=auto.fit_type, + num_spheres=auto.num_spheres, + sphere_density=auto.sphere_density, + surface_radius=auto.surface_radius, + iterations=auto.iterations, + collision_sphere_buffer=auto.collision_sphere_buffer, + device=str(self._curobo_device), + ) + + def _robot_yaml_cache_key( + self, + urdf_path: str, + control_part: str, + tool_frame: str | None, + auto: CuroboAutoGenCfg, + ) -> str: + """Hash the URDF path/content and fit parameters into a stable cache key.""" + hasher = hashlib.md5() + hasher.update(_CUROBO_ROBOT_YAML_GENERATOR_VERSION.encode("utf-8")) + hasher.update(urdf_path.encode("utf-8")) + try: + with open(urdf_path, "rb") as urdf_file: + hasher.update(urdf_file.read()) + except OSError: + pass + hasher.update(control_part.encode("utf-8")) + hasher.update((tool_frame or "").encode("utf-8")) + hasher.update(auto.fit_type.encode("utf-8")) + hasher.update(str(auto.num_spheres).encode("utf-8")) + hasher.update(str(auto.sphere_density).encode("utf-8")) + hasher.update(str(auto.surface_radius).encode("utf-8")) + hasher.update(str(auto.iterations).encode("utf-8")) + hasher.update(str(auto.collision_sphere_buffer).encode("utf-8")) + return hasher.hexdigest() + + def _auto_generate_world_yaml(self, world_cfg: CuroboWorldCfg) -> str: + """Return a cached cuRobo world YAML path generated from ``rigid_objects``. + + Mirrors :meth:`_auto_generate_robot_yaml`: a content-hashed YAML is written + to the cuRobo cache directory (reusing :attr:`CuroboAutoGenCfg.cache_dir`) + on the first plan and reused thereafter. Sphere-fit parameters come from + :class:`CuroboAutoGenCfg` so robot and world fitting are configured together. + """ + from .curobo_yaml import generate_curobo_world_yaml + + rigid_objects = world_cfg.rigid_objects + if not rigid_objects: + logger.log_error( + "_auto_generate_world_yaml requires non-empty rigid_objects.", + ValueError, + ) + assert rigid_objects is not None # log_error raises above; narrows type + auto = self.cfg.auto_gen + cache_dir = auto.cache_dir or os.path.join( + os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), + "embodichain_curobo", + ) + cache_key = self._world_yaml_cache_key(world_cfg) + cache_path = os.path.join(cache_dir, f"world_{cache_key}.yml") + if not auto.force and os.path.exists(cache_path): + logger.log_info(f"cuRobo world YAML cache hit: {cache_path}") + return cache_path + logger.log_info( + f"Auto-generating cuRobo world YAML from {len(rigid_objects)} " + f"RigidObject(s) ({world_cfg.obstacle_representation}) -> {cache_path}" + ) + return generate_curobo_world_yaml( + rigid_objects, + cache_path, + representation=world_cfg.obstacle_representation, + fit_type=auto.fit_type, + num_spheres=auto.num_spheres, + sphere_density=auto.sphere_density, + surface_radius=auto.surface_radius, + iterations=auto.iterations, + collision_sphere_buffer=auto.collision_sphere_buffer, + device=str(self._curobo_device), + ) + + def _world_yaml_cache_key(self, world_cfg: CuroboWorldCfg) -> str: + """Hash per-object mesh/pose + representation + fit params into a cache key. + + Includes each object's vertex/face/pose bytes so editing the simulator + geometry or moving a static obstacle regenerates the YAML, matching the + robot-YAML cache's URDF-content inclusion. + """ + hasher = hashlib.md5() + hasher.update(world_cfg.obstacle_representation.encode("utf-8")) + auto = self.cfg.auto_gen + hasher.update(auto.fit_type.encode("utf-8")) + hasher.update(str(auto.num_spheres).encode("utf-8")) + hasher.update(str(auto.sphere_density).encode("utf-8")) + hasher.update(str(auto.surface_radius).encode("utf-8")) + hasher.update(str(auto.iterations).encode("utf-8")) + hasher.update(str(auto.collision_sphere_buffer).encode("utf-8")) + for idx, obj in enumerate(world_cfg.rigid_objects or []): + name = getattr(obj, "uid", None) or f"obstacle_{idx}" + hasher.update(name.encode("utf-8")) + vertices = obj.get_vertices(env_ids=[0], scale=True)[0] + faces = obj.get_triangles(env_ids=[0])[0] + pose = obj.get_local_pose(to_matrix=False)[0] + hasher.update( + vertices.detach().to("cpu").to(torch.float32).numpy().tobytes() + ) + hasher.update(faces.detach().to("cpu").numpy().tobytes()) + hasher.update(pose.detach().to("cpu").to(torch.float32).numpy().tobytes()) + return hasher.hexdigest() + + 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]) + + # ------------------------------------------------------------------ + # Segment planning + # ------------------------------------------------------------------ + + def _plan_segments( + self, + target_states: list[PlanState], + start: torch.Tensor, + backends: dict[MoveType, "_CuroboBackend"], + options: CuroboPlanOptions, + sim_base_pose_inv: torch.Tensor | None = None, + ) -> PlanResult: + """Plan each waypoint segment sequentially and assemble a PlanResult. + + Each segment's goal is converted to the cuRobo base frame and solved by + the cached in-process V2 planner. Segment extraction, planning-time + budget checks, junction de-duplication, and rectangular assembly all + stay on cuRobo's CUDA device. The assembled result is copied to the + simulation device once at the API boundary. + """ + 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._curobo_device) + current = start.clone() + + for seg_idx, target in enumerate(target_states): + self._validate_segment_batch(target, B, seg_idx) + backend = backends[target.move_type] + 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, sim_base_pose_inv + ) + start_time = time.time() + capture_mode = ( + _torch_cuda_graph_capture_mode( + self.cfg.cuda_graph_capture_error_mode + ) + if backend.use_cuda_graph + else nullcontext() + ) + with capture_mode, torch.cuda.device(self._curobo_device): + v2_result = backend.planner.plan_pose( + goal, current_state, max_attempts=max_attempts + ) + logger.log_info( + f"cuRobo plan_pose segment {seg_idx} cost time: " + f"{time.time() - start_time:.4f}s" + ) + 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) + start_time = time.time() + capture_mode = ( + _torch_cuda_graph_capture_mode( + self.cfg.cuda_graph_capture_error_mode + ) + if backend.use_cuda_graph + else nullcontext() + ) + with capture_mode, torch.cuda.device(self._curobo_device): + v2_result = backend.planner.plan_cspace( + goal_state, current_state, max_attempts=max_attempts + ) + logger.log_info( + f"cuRobo plan_cspace segment {seg_idx} cost time: " + f"{time.time() - start_time:.4f}s" + ) + else: + logger.log_error( + f"cuRobo does not support move_type {target.move_type}.", + ValueError, + ) + + 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._curobo_device + ) + seg_positions = current.unsqueeze(1) + seg_dt = torch.zeros( + B, 1, dtype=torch.float32, device=self._curobo_device + ) + else: + seg_success, seg_positions, seg_dt = self._extract_segment( + v2_result, backend + ) + seg_success = seg_success.to(self._curobo_device) & alive + 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) + + 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 _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]: + """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._curobo_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, D = position.shape + # Compute the per-env valid length once. This scalar extraction is the + # only synchronization needed before rectangular trajectory assembly. + max_len = max(int((last_tstep + 1).max().item()), 1) + cap = min(max_len, T) + lengths = (last_tstep + 1).clamp(min=1, max=cap).long().to(self._curobo_device) + # A single gather both trims to each env's length and pads by repeating + # the last valid sample: src[b, t] = t if t < length[b] else length[b] - 1. + # cap <= T guarantees src < T, so the gather never indexes out of bounds. + position = position.float().to(self._curobo_device) + arange = torch.arange(max_len, device=self._curobo_device) + src = torch.where( + arange[None, :] < lengths[:, None], + arange[None, :], + lengths[:, None] - 1, + ).long() + full = position.gather(1, src.unsqueeze(-1).expand(-1, -1, D)) + + seg_positions = self._map_curobo_to_sim(full, traj.joint_names, backend) + seg_dt = self._extract_dt(traj, lengths, max_len, B) + return success, seg_positions, seg_dt + + def _map_curobo_to_sim( + self, + full_positions: torch.Tensor, + curobo_joint_names: list[str], + backend: "_CuroboBackend", + ) -> torch.Tensor: + """Map a full cuRobo trajectory to simulator control-part joint order. + + The cuRobo joint order is fixed for a planner's life, so the column + gather index is built once (cached on ``backend``) and reused instead of + recomputing O(D^2) ``.index()`` lookups on every segment. + """ + sig = tuple(curobo_joint_names) + if ( + backend.curobo_joint_names_sig != sig + or backend.curobo_to_sim_col_idx is None + ): + sim_to_curobo = backend.profile.sim_to_curobo_joint_names + cols: list[int] = [] + 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( + 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)) + backend.curobo_to_sim_col_idx = torch.as_tensor( + cols, dtype=torch.long, device=self._curobo_device + ) + backend.curobo_joint_names_sig = sig + return full_positions[..., backend.curobo_to_sim_col_idx].to( + dtype=torch.float32 + ) + + def _extract_dt( + self, + traj: "Any", + lengths: torch.Tensor, + max_len: int, + B: int, + ) -> torch.Tensor: + """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. ``lengths`` + is the per-env valid-length tensor (computed once in + :meth:`_extract_segment` and reused here) so this is a vectorized mask + instead of a per-env Python loop with ``.item()`` syncs. + """ + raw_dt = getattr(traj, "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: + dt = torch.full( + (B, 1), + float(self.cfg.interpolation_dt), + device=self._curobo_device, + dtype=torch.float32, + ) + 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._curobo_device, dtype=torch.float32) + if dt.shape[-1] == 1: + # Scalar dt per env: out[b, t] = interval[b] for 1 <= t < length[b], + # else 0 - one vectorized mask multiply (was a per-env Python loop). + interval = dt[:, 0].to(self._curobo_device, dtype=torch.float32) + arange = torch.arange(max_len, device=self._curobo_device) + mask = (arange[None, :] >= 1) & (arange[None, :] < lengths[:, None]) + return interval[:, None] * mask + + # 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._curobo_device, dtype=torch.float32) + 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._curobo_device) + if tt.dim() == 2: + tt = tt.squeeze(-1) + return tt[:B].to(self._curobo_device) + return torch.full((B,), float(tt), device=self._curobo_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.""" + # One D2H sync for the whole batch (was B per-env `if alive[b]:` syncs, + # each forcing the GPU pipeline to drain). The rest of the loop reads + # Python bools and GPU tensors whose .shape / .cat do not sync. + alive_list = alive.tolist() + env_lengths: list[int] = [] + for b in range(B): + if alive_list[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._curobo_device, dtype=torch.float32 + ) + dt = torch.zeros(B, max_len, device=self._curobo_device, dtype=torch.float32) + for b in range(B): + if alive_list[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.to(self.device), + positions=positions.to(self.device), + dt=dt.to(self.device), + duration=duration.to(self.device), + ) + + # ------------------------------------------------------------------ + # cuRobo state / goal construction + # ------------------------------------------------------------------ + + def _to_curobo_joint_state( + self, current: torch.Tensor, backend: "_CuroboBackend" + ) -> "Any": + """Build a cuRobo ``JointState`` from simulator-order joint positions.""" + 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, + ) + curobo_names = list(backend.planner.joint_names) + if backend.sim_to_curobo_col_idx is None: + curobo_to_sim = { + backend.profile.sim_to_curobo_joint_names[sim_name]: idx + for idx, sim_name in enumerate(backend.sim_joint_names) + } + backend.sim_to_curobo_col_idx = torch.as_tensor( + [curobo_to_sim[name] for name in curobo_names], + dtype=torch.long, + device=self._curobo_device, + ) + position = current.to(self._curobo_device, dtype=torch.float32).index_select( + -1, backend.sim_to_curobo_col_idx + ) + return self._bindings.JointState.from_position( + position, joint_names=curobo_names + ) + + def _to_curobo_pose_goal( + self, + xpos: torch.Tensor, + backend: "_CuroboBackend", + sim_base_pose_inv: torch.Tensor | None = None, + ) -> "Any": + """Build a cuRobo pose goal from a simulator-world TCP pose.""" + goal_matrix = self._to_curobo_base_tool_matrix(xpos, backend, sim_base_pose_inv) + position, quaternion = _matrix_to_position_quaternion(goal_matrix) + 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 from simulator-order joint positions.""" + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self._curobo_device) + if qpos.dim() == 1: + qpos = qpos.unsqueeze(0) + return self._to_curobo_joint_state(qpos, backend) + + def _tcp_to_tool_pose( + self, tcp_pose: torch.Tensor, backend: "_CuroboBackend" + ) -> 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, + ) + tool_to_frame = self._tool_to_frame_matrix(backend) + if tool_to_frame is None: + return tcp_pose + return tcp_pose @ tool_to_frame + + def _tool_to_frame_matrix(self, backend: "_CuroboBackend") -> torch.Tensor | None: + """Cached inverse of the profile's fixed tool_frame->TCP transform. + + ``None`` means the tool frame is already the TCP (the common auto-derived + case). Built once per backend and reused across plans instead of calling + ``torch.linalg.inv`` on every EEF segment. + """ + if backend.tool_to_frame_matrix is not None: + return backend.tool_to_frame_matrix + profile = backend.profile + if profile.tool_frame_to_tcp is None: + return None + frame_to_tcp = torch.as_tensor( + profile.tool_frame_to_tcp, + dtype=torch.float32, + device=self._curobo_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, + ) + backend.tool_to_frame_matrix = pose_inv(frame_to_tcp) + return backend.tool_to_frame_matrix + + def _sim_world_to_curobo_base_pose( + self, + world_pose: torch.Tensor, + backend: "_CuroboBackend", + sim_base_pose_inv: torch.Tensor | None = None, + ) -> 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. + + ``sim_base_pose_inv`` is the precomputed inverse of the live sim base + pose; the :meth:`plan` hot path passes it so K segments and N dynamic + obstacles reuse one inverse (the robot does not move during planning). + The public path leaves it ``None`` and computes it here via + :func:`pose_inv` (closed-form, cheaper and more stable than + ``torch.linalg.inv``). + """ + 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] + if sim_base_pose_inv is None: + sim_base_pose = self._get_sim_base_pose(backend, batch_size) + sim_base_pose_inv = pose_inv(sim_base_pose) + sim_base_to_curobo = self._sim_base_to_curobo_matrix(backend).expand( + batch_size, -1, -1 + ) + return torch.bmm( + sim_base_to_curobo, + torch.bmm(sim_base_pose_inv, world_pose), + ) + + def _sim_base_to_curobo_matrix(self, backend: "_CuroboBackend") -> torch.Tensor: + """Cached fixed sim-base -> cuRobo-base transform (eye when ``None``). + + Built once per backend and reused across plans instead of + ``torch.as_tensor``-ing the profile list on every call. + """ + if backend.sim_base_to_curobo_base_matrix is not None: + return backend.sim_base_to_curobo_base_matrix + profile_transform = backend.profile.sim_base_to_curobo_base + if profile_transform is None: + matrix = torch.eye(4, dtype=torch.float32, device=self._curobo_device) + else: + matrix = torch.as_tensor( + profile_transform, + dtype=torch.float32, + device=self._curobo_device, + ) + if matrix.shape != (4, 4): + logger.log_error( + "sim_base_to_curobo_base must be a homogeneous (4, 4) " + f"transform, got {tuple(matrix.shape)}.", + ValueError, + ) + backend.sim_base_to_curobo_base_matrix = matrix + return matrix + + 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 a solver with " + "root_link_name for cuRobo world-frame conversion.", + ValueError, + ) + assert root_link_name is not None # log_error raises above; narrows type + 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._curobo_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 + # ------------------------------------------------------------------ + + def update_dynamic_obstacles( + self, + poses: dict[str, torch.Tensor] | None, + backend: "_CuroboBackend | None" = None, + sim_base_pose_inv: torch.Tensor | None = None, + ) -> None: + """Update named dynamic obstacle poses on cached cuRobo collision worlds. + + Args: + poses: Mapping of obstacle name -> ``(B, 4, 4)`` world pose. ``None`` + is a no-op. + backend: Specific cached backend to update. If ``None``, updates all + cached backends. + sim_base_pose_inv: Precomputed inverse of the live sim base pose for + ``backend``'s batch size, reused across all obstacles. Only + consulted when its batch matches the obstacle pose batch. + """ + 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()) + ) + if backend is None and self.cfg.world.multi_env: + batch_sizes = {cached.batch_size for cached in backends} + if len(batch_sizes) > 1: + logger.log_error( + "Cannot update all cached multi-env cuRobo backends with " + "different batch sizes. Pass the intended backend explicitly.", + ValueError, + ) + + inv_cache: dict[int, torch.Tensor] = {} + for name, pose_tensor in poses.items(): + pose_tensor = torch.as_tensor( + pose_tensor, device=self._curobo_device, dtype=torch.float32 + ) + b = pose_tensor.shape[0] + for cached_backend in backends: + key = id(cached_backend) + inv = inv_cache.get(key) + if inv is None or inv.shape[0] != b: + if ( + backend is not None + and sim_base_pose_inv is not None + and sim_base_pose_inv.shape[0] == b + ): + inv = sim_base_pose_inv + else: + inv = pose_inv(self._get_sim_base_pose(cached_backend, b)) + inv_cache[key] = inv + curobo_pose = self._sim_world_to_curobo_base_pose( + pose_tensor, cached_backend, inv + ) + self._update_backend_obstacle(name, curobo_pose, cached_backend) + + def _update_backend_obstacle( + self, name: str, pose_tensor: torch.Tensor, backend: "_CuroboBackend" + ) -> None: + """Apply one 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 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=positions[env_idx], quaternion=quaternions[env_idx] + ) + backend.planner.scene_collision_checker.update_obstacle_pose( + name, pose, env_idx=env_idx + ) + return + + 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 + ) + + # ------------------------------------------------------------------ + # In-process goal conversion and lifecycle + # ------------------------------------------------------------------ + + def _to_curobo_base_tool_matrix( + self, + xpos: torch.Tensor, + backend: "_CuroboBackend", + sim_base_pose_inv: torch.Tensor | None = None, + ) -> torch.Tensor: + """Convert a batched sim-world TCP pose to a cuRobo-base tool-frame matrix. + + Pure-tensor composition of :meth:`_sim_world_to_curobo_base_pose` and + :meth:`_tcp_to_tool_pose`. + ``sim_base_pose_inv`` is the per-plan cached base-pose inverse (see + :meth:`plan`). + """ + xpos = torch.as_tensor(xpos, device=self._curobo_device, dtype=torch.float32) + xpos = self._sim_world_to_curobo_base_pose(xpos, backend, sim_base_pose_inv) + xpos = self._tcp_to_tool_pose(xpos, backend) + return xpos + + def close(self) -> None: + """Destroy every cached in-process cuRobo planner.""" + for backend in list(self._backend_cache.values()): + self._close_planner(backend.planner) + self._backend_cache.clear() + + def __del__(self) -> None: # pragma: no cover - best-effort GC cleanup + try: + self.close() + except Exception: + pass + + +@dataclass +class _CuroboBackend: + """Cached in-process V2 planner and its EmbodiChain-side metadata.""" + + planner: "Any" + control_part: str + sim_joint_names: list[str] + tool_frame: str + profile: _CuroboProfile + batch_size: int + use_cuda_graph: bool + planning_mode: MoveType + # Lazily-built device-tensor caches for the shared post-processing. The + # cuRobo joint order and the profile's fixed transforms are stable for a + # planner's life, so these are built once on first use and reused across + # plans instead of recomputing per segment / per plan. + sim_to_curobo_col_idx: torch.Tensor | None = None + curobo_to_sim_col_idx: torch.Tensor | None = None + curobo_joint_names_sig: tuple[str, ...] | None = None + tool_to_frame_matrix: torch.Tensor | None = None + sim_base_to_curobo_base_matrix: torch.Tensor | None = None diff --git a/embodichain/lab/sim/planners/curobo/curobo_yaml.py b/embodichain/lab/sim/planners/curobo/curobo_yaml.py new file mode 100644 index 000000000..1b24eec68 --- /dev/null +++ b/embodichain/lab/sim/planners/curobo/curobo_yaml.py @@ -0,0 +1,646 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Generate cuRobo V2 configuration YAMLs from EmbodiChain simulator objects. + +The :func:`generate_curobo_robot_yaml` helper pulls the robot's URDF path and +each link's collision mesh (vertices/faces) from the simulator, fits collision +spheres to every link mesh with cuRobo's sphere-fitting library, and writes a +complete cuRobo V2 robot configuration YAML. The cuRobo planner adapter calls +this automatically (with on-disk caching) on the first plan; see +:class:`~embodichain.lab.sim.planners.curobo.curobo_planner.CuroboAutoGenCfg`. + +:func:`generate_curobo_world_yaml` builds the cuRobo collision-world YAML from +live :class:`~embodichain.lab.sim.objects.RigidObject` meshes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +import torch + +from embodichain.utils import logger +from embodichain.utils.math import matrix_from_quat, quat_from_matrix + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import RigidObject, Robot + +__all__ = ["generate_curobo_robot_yaml", "generate_curobo_world_yaml"] + + +def _parse_mimic_joint_names(urdf_path: str) -> set[str]: + """Return the names of URDF joints that mimic another joint. + + cuRobo's URDF parser folds each ```` joint into its active joint: the + mimic joint's body takes the active joint's name, so the mimic joint has no + independent entry in the kinematics tree. cuRobo therefore rejects mimic + joints in ``cspace`` and ``lock_joints`` - locking one raises ``KeyError`` + because cuRobo finds no body for it. They must be excluded from both, so + cuRobo drives them from their active joint instead. + + cuRobo's ``UrdfRobotParser`` exposes no public accessor for mimic joints (a + prior ``get_mimic_joint_map`` call did not exist on this cuRobo version and + silently left the set empty), so they are read directly from the URDF XML - + the same ```` tags cuRobo itself parses. + + Args: + urdf_path: Path to the robot URDF file. + + Returns: + The set of joint names declared with a ```` child element. Empty + if the URDF cannot be parsed (a warning is logged). + """ + import xml.etree.ElementTree as ET + + mimic_joints: set[str] = set() + try: + root = ET.parse(urdf_path).getroot() + except Exception as exc: # noqa: BLE001 + logger.log_warning(f"Could not parse mimic joints from URDF ({exc}).") + return mimic_joints + for joint in root.findall("joint"): + if joint.find("mimic") is not None: + name = joint.get("name") + if name is not None: + mimic_joints.add(name) + return mimic_joints + + +def generate_curobo_robot_yaml( + robot: Robot, + control_part: str, + output_path: str, + *, + tool_frame: str | None = None, + urdf_path: str | None = None, + fit_type: str = "morphit", + num_spheres: int | None = None, + sphere_density: float = 1.0, + surface_radius: float = 0.005, + iterations: int = 200, + collision_sphere_buffer: float = 0.0, + max_acceleration: float = 15.0, + max_jerk: float = 500.0, + device: str = "cuda:0", +) -> str: + """Fit collision spheres to each robot link's mesh and write a cuRobo robot YAML. + + Extracts the URDF path and per-link vertices/faces from ``robot``, fits + collision spheres to every link mesh with cuRobo's :func:`fit_spheres_to_mesh`, + and writes a complete cuRobo V2 robot configuration YAML that the cuRobo + planner loads as its robot model. + + .. attention:: + Requires a CUDA GPU and cuRobo installed (sphere fitting runs on GPU). + Link meshes from ``robot.get_link_vert_face`` are assumed to be in the + link-local rest frame -- the convention cuRobo collision spheres use, + since cuRobo applies each link's transform via FK at runtime. + + Args: + robot: The EmbodiChain robot to generate a config for. + control_part: Control-part name whose joints stay active; every other + actuated joint is pinned via ``lock_joints``. + output_path: Destination YAML file path. + tool_frame: cuRobo tool frame (a URDF link name) to plan to. If ``None``, + defaults to the last link of the control part. + urdf_path: URDF to generate the cuRobo model from. If ``None`` (default), + uses ``robot.cfg.fpath`` -- the *assembled* URDF that includes every + mounted component (arm + gripper). Pass this explicitly when the + caller already resolved the URDF (the planner does, so the on-disk + cache key and the generation use the same file). Must be the full + assembled URDF, not a solver's sub-chain URDF, or gripper links are + silently dropped from the collision model. + fit_type: cuRobo sphere-fit strategy - ``"morphit"`` (default, best), + ``"voxel"`` (faster), or ``"surface"`` (crude, fixed radius). + num_spheres: Per-link sphere count. If ``None``, cuRobo auto-estimates + it from the link's bounding-box volume. + sphere_density: Multiplier on the auto sphere count (ignored when + ``num_spheres`` is set). + surface_radius: Fixed radius used only by the ``surface`` strategy. + iterations: Adam iterations for the ``morphit`` strategy. + collision_sphere_buffer: Padding added to every sphere's radius (m). + max_acceleration: cspace maximum acceleration. + max_jerk: cspace maximum jerk. + device: CUDA device for sphere fitting. + + Returns: + The ``output_path`` that was written. + + Raises: + ImportError: If cuRobo or trimesh is not installed. + RuntimeError: If CUDA is unavailable or no spheres could be fitted. + """ + import os + + import trimesh + import yaml + + from curobo._src.geom.sphere_fit.fit_spheres import fit_spheres_to_mesh + from curobo._src.geom.sphere_fit.types import SphereFitType + from curobo._src.robot.parser.parser_urdf import UrdfRobotParser + from curobo.types import DeviceCfg + + if not torch.cuda.is_available(): + raise RuntimeError("generate_curobo_robot_yaml requires a CUDA GPU.") + fit_type_map = { + "morphit": SphereFitType.MORPHIT, + "voxel": SphereFitType.VOXEL, + "surface": SphereFitType.SURFACE, + } + if fit_type not in fit_type_map: + raise ValueError( + f"fit_type must be one of {list(fit_type_map)}, got {fit_type!r}." + ) + fit_type_enum = fit_type_map[fit_type] + device_cfg = DeviceCfg(device=device) + + urdf_path = urdf_path or robot.cfg.fpath + link_vert_dict: dict = {} + link_face_dict: dict = {} + for link_name in robot.get_link_names() or []: + verts, faces = robot.get_link_vert_face(link_name) + link_vert_dict[link_name] = verts + link_face_dict[link_name] = faces + + # 1. Parse the URDF kinematic tree (no meshes) for base_link + parent map. + # ``robot.root_link_name`` is avoided because it touches an uninitialized + # ``entities`` attribute on some Robot instances; cuRobo's parser resolves + # the root link directly from the URDF. + # Mimic joints are detected from the URDF XML (not cuRobo's parser, which + # exposes no mimic accessor) so they can be excluded from cspace/lock_joints + # in step 4/5 - cuRobo folds them into their active joint and raises + # KeyError if they are locked. + mimic_joints: set[str] = _parse_mimic_joint_names(urdf_path) + base_link: str | None = None + urdf_parent_map: dict[str, str | None] = {} + try: + parser = UrdfRobotParser(urdf_path, load_meshes=False, build_scene_graph=True) + parser.build_link_parent() + base_link = parser.root_link + # Build the full parent map for every URDF link so self_collision_ignore + # can walk multiple hops (the parent of a non-collision link still + # connects two collision links, e.g. fr3_link8 between fr3_link7 and + # fr3_hand). + for link_name in parser.get_link_names_from_urdf(): + try: + urdf_parent_map[link_name] = parser.get_link_parameters( + link_name + ).parent_link_name + except Exception: # noqa: BLE001 (e.g. root link has no parent entry) + urdf_parent_map[link_name] = None + except Exception as exc: # noqa: BLE001 + logger.log_warning(f"Could not parse URDF kinematic tree ({exc}).") + if base_link is None: + link_names_fb = robot.get_link_names() or [] + base_link = getattr(robot.cfg, "base_link_name", None) or ( + link_names_fb[0] if link_names_fb else "base_link" + ) + + # 2. Fit collision spheres per link from the simulator meshes. + collision_spheres: dict[str, list[dict]] = {} + for link_name, verts in link_vert_dict.items(): + faces = link_face_dict[link_name] + if verts is None or faces is None or verts.numel() == 0 or faces.numel() == 0: + continue + verts_np = torch.as_tensor(verts).detach().to(torch.float32).cpu().numpy() + faces_np = torch.as_tensor(faces).detach().to(torch.int64).cpu().numpy() + mesh = trimesh.Trimesh(vertices=verts_np, faces=faces_np, process=False) + if len(mesh.vertices) == 0: + continue + mesh.fill_holes() + trimesh.repair.fix_normals(mesh) + trimesh.repair.fix_inversion(mesh) + trimesh.repair.fix_winding(mesh) + try: + fit_result = fit_spheres_to_mesh( + mesh, + num_spheres=num_spheres, + sphere_density=sphere_density, + surface_radius=surface_radius, + fit_type=fit_type_enum, + iterations=iterations, + device_cfg=device_cfg, + ) + except Exception as exc: # noqa: BLE001 + logger.log_warning(f"Sphere fitting failed for link {link_name!r}: {exc}") + continue + if fit_result.num_spheres == 0: + continue + collision_spheres[link_name] = [ + {"center": list(c), "radius": float(r)} + for c, r in zip( + fit_result.centers.detach().cpu().tolist(), + fit_result.radii.detach().cpu().tolist(), + ) + ] + + if not collision_spheres: + raise RuntimeError( + "No collision spheres could be fitted from the robot's link meshes." + ) + collision_link_names = list(collision_spheres.keys()) + + # 3. self_collision_ignore: ignore link pairs within two kinematic hops + # (parent/grandparent, children/grandchildren, siblings). cuRobo's curated + # profiles (e.g. franka.yml) ignore adjacent-plus-near links because their + # spheres physically overlap near joints; a neighbor-only matrix leaves + # those pairs colliding and makes reachable start poses fail validation. + self_collision_ignore: dict[str, list[str]] = {} + if urdf_parent_map: + children_map: dict[str, list[str]] = {} + for link_name, parent in urdf_parent_map.items(): + if parent is not None: + children_map.setdefault(parent, []).append(link_name) + collision_set = set(collision_link_names) + + def _two_hop_neighbors(link: str) -> set[str]: + neighbors: set[str] = set() + parent = urdf_parent_map.get(link) + if parent is not None: + neighbors.add(parent) + grandparent = urdf_parent_map.get(parent) + if grandparent is not None: + neighbors.add(grandparent) + for sibling in children_map.get(parent, []): + if sibling != link: + neighbors.add(sibling) + for child in children_map.get(link, []): + neighbors.add(child) + for grandchild in children_map.get(child, []): + neighbors.add(grandchild) + return neighbors + + for link_name in collision_link_names: + self_collision_ignore[link_name] = [ + n for n in _two_hop_neighbors(link_name) if n in collision_set + ] + + # 4. cspace from the robot's joints + init qpos. Mimic joints are excluded - + # cuRobo drives them from their active joint and rejects them in cspace. + joint_names = list(robot.joint_names) + init_qpos = list(robot.cfg.init_qpos) if robot.cfg.init_qpos is not None else [] + if len(init_qpos) != len(joint_names): + logger.log_warning( + "init_qpos length does not match joint_names; using current qpos." + ) + try: + init_qpos = robot.get_qpos()[0].detach().cpu().tolist() + except Exception: # noqa: BLE001 + init_qpos = [0.0] * len(joint_names) + cspace_pairs = [ + (jname, float(val)) + for jname, val in zip(joint_names, init_qpos) + if jname not in mimic_joints + ] + cspace = { + "joint_names": [j for j, _ in cspace_pairs], + "default_joint_position": [v for _, v in cspace_pairs], + "max_acceleration": float(max_acceleration), + "max_jerk": float(max_jerk), + "cspace_distance_weight": [1.0] * len(cspace_pairs), + "null_space_weight": [1.0] * len(cspace_pairs), + } + + # 5. lock_joints: actuated joints outside the control part, pinned to init values. + # Mimic joints are already excluded from cspace_pairs (see step 4). + control_joints = set((robot.control_parts or {}).get(control_part, [])) + lock_joints: dict[str, float] = { + jname: val for jname, val in cspace_pairs if jname not in control_joints + } + + # 6. tool_frames default to the last link of the control part. + if tool_frame is None: + part_links = robot.get_control_part_link_names(control_part) + if not part_links: + raise RuntimeError( + f"Control part {control_part!r} has no links; specify tool_frame." + ) + tool_frame = part_links[-1] + + # 7. Assemble and write the YAML, mirroring franka.yml's schema. + data = { + "robot_cfg": { + "kinematics": { + "format_version": 2.0, + "base_link": base_link, + "urdf_path": urdf_path, + "asset_root_path": os.path.dirname(urdf_path), + "tool_frames": [tool_frame], + "collision_link_names": collision_link_names, + "collision_spheres": collision_spheres, + "collision_sphere_buffer": float(collision_sphere_buffer), + "mesh_link_names": collision_link_names, + "self_collision_buffer": {ln: 0.0 for ln in collision_link_names}, + "self_collision_ignore": self_collision_ignore, + "lock_joints": lock_joints, + "cspace": cspace, + "use_global_cumul": True, + } + } + } + + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + with open(output_path, "w") as yaml_file: + yaml.dump(data, yaml_file, default_flow_style=False, sort_keys=False) + return output_path + + +# ============================================================================= +# World (obstacle) YAML generation from RigidObject meshes +# ============================================================================= + + +_REPRESENTATIONS = ("cuboid", "mesh", "sphere") + + +def _mesh_to_obstacle_entry( + name: str, + vertices: torch.Tensor, + faces: torch.Tensor, + pose: torch.Tensor, + *, + representation: str = "cuboid", + fit_type: str = "voxel", + num_spheres: int | None = None, + sphere_density: float = 1.0, + surface_radius: float = 0.005, + iterations: int = 200, + collision_sphere_buffer: float = 0.0, + device: str = "cuda:0", +) -> list[tuple[str, str, dict]]: + """Convert one mesh + pose into cuRobo world-YAML obstacle entry/entries. + + Pure tensor helper (no simulator / cuRobo import for ``cuboid``/``mesh``) so + it is unit-testable without CUDA. ``sphere`` lazily imports cuRobo + trimesh + and runs on CUDA. + + Args: + name: Obstacle name (cuRobo key under ``cuboid``/``mesh``/``sphere``). + vertices: Mesh vertices ``(V, 3)`` in the object's local frame. + faces: Triangle indices ``(F, 3)`` (any integer dtype). + pose: Object pose as ``(x, y, z, qw, qx, qy, qz)`` ``(7,)`` or a + homogeneous ``(4, 4)`` matrix, expressed in the cuRobo world/base + frame (the same frame static collision YAMLs are authored in). + representation: ``"cuboid"`` (local-frame AABB -> OBB via ``pose``, + default), ``"mesh"`` (exact triangle mesh), or ``"sphere"`` (fit + spheres with cuRobo's :func:`fit_spheres_to_mesh`). + fit_type: cuRobo sphere-fit strategy (``"voxel"``/``"morphit"``/ + ``"surface"``); only used by ``"sphere"``. + num_spheres: Per-mesh sphere count; ``None`` auto-estimates (sphere only). + sphere_density: Multiplier on the auto sphere count (sphere only). + surface_radius: Fixed radius for the ``"surface"`` strategy (sphere only). + iterations: Adam iterations for ``"morphit"`` (sphere only). + collision_sphere_buffer: Padding added to each fitted radius (sphere only). + device: CUDA device for sphere fitting (sphere only). + + Returns: + A list of ``(top_level_key, obstacle_name, fields)`` tuples. ``cuboid``/ + ``mesh`` return one entry; ``sphere`` returns one entry per fitted sphere. + + Raises: + ValueError: If ``representation`` is unsupported, ``pose`` is malformed, + or the mesh has no geometry for the requested representation. + RuntimeError: If ``"sphere"`` is requested without CUDA. + ImportError: If ``"sphere"`` is requested without cuRobo/trimesh. + """ + if representation not in _REPRESENTATIONS: + raise ValueError( + f"representation must be one of {_REPRESENTATIONS}, got {representation!r}." + ) + + vertices = ( + torch.as_tensor(vertices, dtype=torch.float32).detach().to("cpu").reshape(-1, 3) + ) + faces = torch.as_tensor(faces).detach().to("cpu") + pose = torch.as_tensor(pose, dtype=torch.float32).detach().to("cpu") + if pose.shape == (4, 4): + position = pose[:3, 3] + quaternion = quat_from_matrix(pose[:3, :3]) # wxyz + pose = torch.cat([position, quaternion]) + if pose.shape != (7,): + raise ValueError( + f"pose must be (7,) [x,y,z,qw,qx,qy,qz] or (4, 4), got {tuple(pose.shape)}." + ) + + if representation == "mesh": + if vertices.numel() == 0 or faces.numel() == 0: + raise ValueError( + f"object {name!r} has no mesh geometry for the 'mesh' representation." + ) + return [ + ( + "mesh", + name, + { + "vertices": vertices.tolist(), + "faces": faces.reshape(-1).to(torch.int64).tolist(), + "pose": pose.tolist(), + }, + ) + ] + + if representation == "cuboid": + if vertices.numel() == 0: + raise ValueError( + f"object {name!r} has no vertices for the 'cuboid' representation." + ) + # Local-frame AABB, emitted as an OBB via the object pose: cuRobo's + # Cuboid is centered at ``pose[:3]`` with ``dims`` along the pose axes. + vmin = vertices.amin(dim=0) + vmax = vertices.amax(dim=0) + dims = vmax - vmin + center_local = (vmin + vmax) / 2.0 + rotation = matrix_from_quat(pose[3:7]) # (3, 3), wxyz + center_world = rotation @ center_local + pose[:3] + cuboid_pose = torch.cat([center_world, pose[3:7]]) + return [("cuboid", name, {"dims": dims.tolist(), "pose": cuboid_pose.tolist()})] + + # representation == "sphere": fit spheres in the local frame, then transform + # centers into the cuRobo world/base frame (Sphere obstacles have no pose/FK). + if vertices.numel() == 0 or faces.numel() == 0: + raise ValueError( + f"object {name!r} has no mesh geometry for the 'sphere' representation." + ) + if not torch.cuda.is_available(): + raise RuntimeError( + "The 'sphere' representation requires CUDA for cuRobo sphere fitting." + ) + + import trimesh + + from curobo._src.geom.sphere_fit.fit_spheres import fit_spheres_to_mesh + from curobo._src.geom.sphere_fit.types import SphereFitType + from curobo.types import DeviceCfg + + fit_type_map = { + "morphit": SphereFitType.MORPHIT, + "voxel": SphereFitType.VOXEL, + "surface": SphereFitType.SURFACE, + } + if fit_type not in fit_type_map: + raise ValueError( + f"fit_type must be one of {list(fit_type_map)}, got {fit_type!r}." + ) + mesh = trimesh.Trimesh( + vertices=vertices.numpy(), + faces=faces.reshape(-1, 3).to(torch.int64).numpy(), + process=False, + ) + mesh.fill_holes() + trimesh.repair.fix_normals(mesh) + trimesh.repair.fix_inversion(mesh) + trimesh.repair.fix_winding(mesh) + fit_result = fit_spheres_to_mesh( + mesh, + num_spheres=num_spheres, + sphere_density=sphere_density, + surface_radius=surface_radius, + fit_type=fit_type_map[fit_type], + iterations=iterations, + device_cfg=DeviceCfg(device=device), + ) + if fit_result.num_spheres == 0: + raise RuntimeError(f"No spheres could be fitted for object {name!r}.") + + centers_local = ( + fit_result.centers.detach().to("cpu").reshape(-1, 3).to(torch.float32) + ) + radii = fit_result.radii.detach().to("cpu").reshape(-1).to(torch.float32) + float( + collision_sphere_buffer + ) + rotation = matrix_from_quat(pose[3:7]) + centers_world = centers_local @ rotation.T + pose[:3] + entries: list[tuple[str, str, dict]] = [] + for i in range(centers_world.shape[0]): + entries.append( + ( + "sphere", + f"{name}_{i}", + { + "position": centers_world[i].tolist(), + "radius": float(radii[i].item()), + }, + ) + ) + return entries + + +def generate_curobo_world_yaml( + rigid_objects: Sequence[RigidObject], + output_path: str, + *, + representation: str = "cuboid", + env_id: int = 0, + fit_type: str = "voxel", + num_spheres: int | None = None, + sphere_density: float = 1.0, + surface_radius: float = 0.005, + iterations: int = 200, + collision_sphere_buffer: float = 0.0, + device: str = "cuda:0", +) -> str: + """Generate a cuRobo V2 scene (world) YAML from a sequence of ``RigidObject``. + + Each object's mesh (``get_vertices`` / ``get_triangles``) and world pose + (``get_local_pose``) are converted into cuRobo obstacle entries under a single + top-level key (``cuboid`` / ``mesh`` / ``sphere``). The cuRobo planner loads + the resulting YAML as its collision world. + + .. attention:: + Poses are written in the cuRobo world/base frame - the same convention as + a hand-authored static collision YAML. When the robot base is offset from + the simulator world origin, rebase the object poses first, or register the + obstacle name in ``CuroboWorldCfg.dynamic_obstacle_names`` and update its + pose at plan time via + :meth:`~embodichain.lab.sim.planners.curobo.curobo_planner.CuroboPlanner.update_dynamic_obstacles`. + + Args: + rigid_objects: ``RigidObject`` instances to bake into the collision world. + output_path: Destination YAML file path. + representation: ``"cuboid"`` (default, AABB->OBB, no CUDA), ``"mesh"`` + (exact triangle mesh, no CUDA), or ``"sphere"`` (cuRobo sphere fit, + requires CUDA + cuRobo + trimesh). + env_id: Environment instance index to read geometry/pose from (the static + world is shared, so env 0 is representative). + fit_type: cuRobo sphere-fit strategy (sphere representation only). + num_spheres: Per-object sphere count; ``None`` auto-estimates (sphere only). + sphere_density: Multiplier on the auto sphere count (sphere only). + surface_radius: Fixed radius for the ``"surface"`` strategy (sphere only). + iterations: Adam iterations for ``"morphit"`` (sphere only). + collision_sphere_buffer: Padding added to each fitted radius (sphere only). + device: CUDA device for sphere fitting (sphere only). + + Returns: + The ``output_path`` that was written. + + Raises: + ValueError: If ``rigid_objects`` is empty or a representation/pose is + invalid. + """ + import os + + import yaml + + rigid_objects = list(rigid_objects) + if not rigid_objects: + raise ValueError("rigid_objects must contain at least one RigidObject.") + + data: dict[str, dict[str, object]] = {} + used_names: set[str] = set() + for idx, obj in enumerate(rigid_objects): + name = getattr(obj, "uid", None) or f"obstacle_{idx}" + if name in used_names: + raise ValueError( + f"Duplicate obstacle name {name!r}; RigidObject uids must be unique." + ) + used_names.add(name) + + vertices = obj.get_vertices(env_ids=[env_id], scale=True)[0] + faces = obj.get_triangles(env_ids=[env_id])[0] + pose = obj.get_local_pose(to_matrix=False)[env_id] + + if vertices is None or faces is None or vertices.numel() == 0: + logger.log_warning( + f"RigidObject {name!r} has no mesh geometry; skipping collision export." + ) + continue + + entries = _mesh_to_obstacle_entry( + name, + vertices, + faces, + pose, + representation=representation, + fit_type=fit_type, + num_spheres=num_spheres, + sphere_density=sphere_density, + surface_radius=surface_radius, + iterations=iterations, + collision_sphere_buffer=collision_sphere_buffer, + device=device, + ) + for top_key, obstacle_name, fields in entries: + data.setdefault(top_key, {})[obstacle_name] = fields + + if not data: + raise ValueError( + "No collision obstacles could be generated from the given RigidObjects." + ) + + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + with open(output_path, "w") as yaml_file: + yaml.dump(data, yaml_file, default_flow_style=False, sort_keys=False) + return output_path diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index d83004957..f9446cff7 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: @@ -156,10 +160,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 +231,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 90e917f9e..9ebdb6cc6 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 5fc897f6c..f328a774f 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. @@ -305,6 +309,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/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py new file mode 100644 index 000000000..5feba1dcd --- /dev/null +++ b/examples/sim/planners/curobo_planner.py @@ -0,0 +1,591 @@ +# ---------------------------------------------------------------------------- +# 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 the CUDA-matched EmbodiChain cuRobo V2 +extra installed in the active environment. Installation instructions: +https://nvlabs.github.io/curobo/latest/getting-started/installation.html +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import torch + +# Prefer the in-repo source over any installed (possibly stale) embodichain +# package, so this example exercises the current code. The demo relies on the +# cuRobo adapter's URDF-based robot-YAML auto-generation, which lives in the +# source tree and may not be present in an older installed copy. +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndEffectorPoseTarget, + MoveEndEffector, + MoveEndEffectorCfg, +) +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.objects import RigidObjectCfg, Robot, RigidObject +from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator +from embodichain.lab.sim.planners.curobo.curobo_planner import ( + CuroboPlannerCfg, + CuroboWorldCfg, +) +import numpy as np +from embodichain.lab.sim.robots import FrankaPandaCfg, URRobotCfg, DexforceW1Cfg +from embodichain.lab.sim.shapes import CubeCfg + +__all__ = ["main"] + + +DEFAULT_RECORD_FPS = 20 +DEFAULT_RECORD_MAX_MEMORY = 2048 +DEFAULT_MAX_ATTEMPTS = 2 +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( + "--max-attempts", + type=int, + default=DEFAULT_MAX_ATTEMPTS, + help=( + "cuRobo planning attempts per request. Lower values are faster; " + "increase this if a harder scene fails to find a path." + ), + ) + 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.", + ) + parser.add_argument( + "--robot", + type=str, + default="franka", + help="Robot type for the cuRobo demo (franka, ur, w1).", + ) + parser.add_argument( + "--sim-device", + choices=("cpu", "cuda"), + default="cpu", + help=( + "Physics device (default: cuda). cuRobo always plans on CUDA even " + "when CPU physics is selected." + ), + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Enable in-process cuRobo CUDA graphs with renderer-compatible " + "thread-local stream capture (default: enabled; use " + "--no-cuda-graph to disable)." + ), + ) + 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. From the EmbodiChain repository root, " + "install the extra matching the CUDA environment: " + '`uv pip install ".[curobo-cu12]"` for CUDA 12.x or ' + '`uv pip install ".[curobo-cu13]"` for CUDA 13.x ' + f"(see {CUROBO_INSTALL_URL})." + ) from exc + + +def _build_scene( + headless: bool, + robot_type: str = "franka", + sim_device: str = "cuda", +) -> tuple[SimulationManager, Robot, RigidObject, torch.Tensor, str]: + """Create the one-environment Franka scene with its shared cuboid.""" + sim = SimulationManager( + SimulationManagerCfg( + headless=headless, + sim_device=sim_device, + num_envs=1, + arena_space=2.0, + ) + ) + if robot_type == "franka": + control_part = "arm" + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict( + { + "uid": "franka", + "robot_type": "panda", + "init_qpos": [0.0, -0.5, 0.0, -2.3, 0.0, 1.8, 0.741, 0.04, 0.04], + } + ) + ) + demo_block_size = [0.18, 0.3, 0.36] + demo_block_position = (0.40, 0.0, 0.18) + + target_xpos = torch.tensor( + [ + [ + [9.9896e-01, 4.3707e-02, -1.2806e-02, 6.5e-01], + [4.3759e-02, -9.9903e-01, 3.7920e-03, 8.5299e-04], + [-1.2628e-02, -4.3484e-03, -9.9991e-01, 2.0e-01], + [0.0000e00, 0.0000e00, 0.0000e00, 1.0000e00], + ] + ], + device=robot.device, + ) + elif robot_type == "ur": + control_part = "arm" + hand_urdf_path = get_data_path( + "BrainCoHandRevo1/BrainCoLeftHand/BrainCoLeftHand.urdf" + ) + hand_attach_xpos = np.eye(4) + try: + from scipy.spatial.transform import Rotation as _Rotation + except ImportError as exc: # pragma: no cover - exercised only without SciPy + raise ImportError( + "The '--robot ur' demo path requires SciPy. Install it with " + "`pip install scipy`." + ) from exc + hand_attach_xpos[:3, :3] = _Rotation.from_rotvec( + [90, 0, 0], degrees=True + ).as_matrix() + robot = sim.add_robot( + cfg=URRobotCfg.from_dict( + { + "robot_type": "ur10", + "uid": "ur10_with_brainco", + "urdf_cfg": { + "components": [ + { + "component_type": "hand", + "urdf_path": hand_urdf_path, + "transform": hand_attach_xpos, + }, + ] + }, + "control_parts": { + "hand": [ + "LEFT_HAND_THUMB1", + "LEFT_HAND_THUMB2", + "LEFT_HAND_INDEX", + "LEFT_HAND_MIDDLE", + "LEFT_HAND_RING", + "LEFT_HAND_PINKY", + ], + }, + "drive_pros": { + "stiffness": {"LEFT_[A-Z|_]+[0-9]?": 1e2}, + "damping": {"LEFT_[A-Z|_]+[0-9]?": 1e1}, + "max_effort": {"LEFT_[A-Z|_]+[0-9]?": 1e3}, + "drive_type": "force", + }, + "solver_cfg": {"arm": {"tcp": np.eye(4)}}, + "init_qpos": [ + 0.0, + -np.pi / 2, + -np.pi / 2, + 2.5, + -np.pi / 2, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.5, + -0.00016, + -0.00010, + -0.00013, + -0.00009, + 0.0, + ], + } + ) + ) + demo_block_size = [0.18, 0.3, 0.36] + demo_block_position = (0.60, 0.0, 0.18) + target_xpos = torch.tensor( + [ + [ + [9.9896e-01, 4.3707e-02, -1.2806e-02, 8.5e-01], + [4.3759e-02, -9.9903e-01, 3.7920e-03, 8.5299e-04], + [-1.2628e-02, -4.3484e-03, -9.9991e-01, 3.0e-01], + [0.0000e00, 0.0000e00, 0.0000e00, 1.0000e00], + ] + ], + device=robot.device, + ) + elif robot_type == "w1": + control_part = "right_arm" + cfg = DexforceW1Cfg.from_dict( + { + "uid": "dexforce_w1", + } + ) + cfg.solver_cfg["left_arm"].tcp = np.array( + [ + [1.0, 0.0, 0.0, 0.012], + [0.0, 1.0, 0.0, 0.04], + [0.0, 0.0, 1.0, 0.11], + [0.0, 0.0, 0.0, 1.0], + ] + ) + cfg.solver_cfg["right_arm"].tcp = np.array( + [ + [1.0, 0.0, 0.0, 0.012], + [0.0, 1.0, 0.0, -0.04], + [0.0, 0.0, 1.0, 0.11], + [0.0, 0.0, 0.0, 1.0], + ] + ) + + cfg.init_qpos = [ + 1.0000e00, + -2.0000e00, + 1.0000e00, + 0.0000e00, + -2.6921e-05, + -2.6514e-03, + -1.5708e00, + 1.4575e00, + -7.8540e-01, + 1.2834e-01, + 1.5708e00, + -2.2310e00, + -7.8540e-01, + 1.4461e00, + -1.5708e00, + 1.6716e00, + 7.8540e-01, + 7.6745e-01, + 0.0000e00, + 3.8108e-01, + 0.0000e00, + 0.0000e00, + 0.0000e00, + 0.0000e00, + 1.5000e00, + 0.0000e00, + 0.0000e00, + 0.0000e00, + 0.0000e00, + 1.5000e00, + 6.9974e-02, + 7.3950e-02, + 6.6574e-02, + 6.0923e-02, + 0.0000e00, + 6.7342e-02, + 7.0862e-02, + 6.3684e-02, + 5.7822e-02, + 0.0000e00, + ] + robot = sim.add_robot(cfg=cfg) + + demo_block_size = [0.2, 0.2, 0.2] + demo_block_position = (0.36, -0.15, 0.88) + target_xpos = torch.tensor( + [ + [ + [2.2020e-03, 3.4217e-01, 9.3964e-01, 4.6395e-01], + [1.5398e-04, -9.3964e-01, 3.4217e-01, -1.7e-01], + [1.0000e00, -6.0877e-04, -2.1218e-03, 6.80e-01], + [0.0000e00, 0.0000e00, 0.0000e00, 1.0000e00], + ] + ], + device=robot.device, + ) + + # robot compute ik success in example + is_success, ik_qpos = robot.compute_ik(pose=target_xpos, name=control_part) + print(f"robot compute ik success: {is_success}, ik_qpos: {ik_qpos}") + + # sim.open_window() + # # sim.update(50) + # current_qpos = robot.get_qpos(name=control_part) + # current_xpos = robot.compute_fk(name=control_part, qpos=current_qpos, to_matrix=True) + # print(f"Current {control_part} TCP pose:\n{current_xpos}") + # import ipdb; ipdb.set_trace() + + else: + raise ValueError(f"Unknown robot type '{robot_type}' for cuRobo demo.") + + if robot is None: + raise RuntimeError(f"Failed to add robot '{robot_type}' to the cuRobo demo.") + # This object is also exported into the cuRobo collision world below via + # CuroboWorldCfg.rigid_objects, so the simulator and planner share geometry + # automatically (no hand-authored collision YAML to keep in sync). + demo_block = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="demo_block", + shape=CubeCfg(size=demo_block_size), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=demo_block_position, + init_rot=(0.0, 0.0, 0.0), + ) + ) + + return sim, robot, demo_block, target_xpos, control_part + + +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 _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, + ) + sim.update(step=step_repeat) + + +def _final_tcp_error(robot: Robot, target: torch.Tensor, control_part: str) -> 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, + ) + # Accept either a single (4, 4) pose or a batched (B, 4, 4) target. + target_pos = target[0, :3, 3] if target.dim() == 3 else target[:3, 3] + return float(torch.linalg.vector_norm(final_pose[0, :3, 3] - target_pos)) + + +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.max_attempts < 1: + raise ValueError("--max-attempts must be at least 1.") + if args.record_fps < 1: + raise ValueError("--record-fps must be at least 1.") + _check_runtime() + sim: SimulationManager | None = None + try: + sim, robot, demo_block, target_xpos, control_part = _build_scene( + args.headless, args.robot, args.sim_device + ) + 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, + world=CuroboWorldCfg(rigid_objects=[demo_block]), + max_attempts=args.max_attempts, + 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 sets the returned trajectory's waypoint count. + # cuRobo's own collision-checked samples are arc-length resampled + # to this count; set CuroboPlannerCfg.preserve_plan_samples=True + # above to keep cuRobo's raw samples (count from interpolation_dt). + sample_interval=30, + ), + ), + name="move_end_effector", + ) + + initial_qpos = robot.get_qpos(name=control_part) + initial_xpos = robot.compute_fk( + qpos=initial_qpos, + name=control_part, + to_matrix=True, + ) + plan_start = time.perf_counter() + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target_xpos))] + ) + 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"[warm-up] 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_xpos, control_part):.4f} m" + ) + + plan_start = time.perf_counter() + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=initial_xpos))] + ) + 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"[Runtime]atomic-action planning duration: {planning_duration:.3f} s") + _replay_full_dof_trajectory( + sim, + robot, + trajectory, + step_repeat=args.step_repeat, + ) + if not args.headless: + input("Press Enter to exit the cuRobo demo...") + finally: + # Guarantee GPU sim resources and the recorder subprocess are released + # even when planning/replay raises mid-demo. + 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/pyproject.toml b/pyproject.toml index b36950503..65163bc88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,22 @@ gensim = [ "bpy", "pyrender==0.1.45" ] +# cuRobo V2 is distributed from its source repository and provides separate +# dependency sets for CUDA 12 and CUDA 13. Keep it optional so CPU-only and +# non-NVIDIA EmbodiChain installations do not pull in CUDA packages. Select +# exactly one of these extras for the target environment. +curobo-cu12 = [ + "nvidia-curobo[cu12] @ git+https://github.com/NVlabs/curobo.git@v0.8.0" +] +curobo-cu12-torch = [ + "nvidia-curobo[cu12-torch] @ git+https://github.com/NVlabs/curobo.git@v0.8.0" +] +curobo-cu13 = [ + "nvidia-curobo[cu13] @ git+https://github.com/NVlabs/curobo.git@v0.8.0" +] +curobo-cu13-torch = [ + "nvidia-curobo[cu13-torch] @ git+https://github.com/NVlabs/curobo.git@v0.8.0" +] [tool.uv.sources] bpy = { index = "blender" } diff --git a/scripts/benchmark/planners/benchmark_curobo_extraction.py b/scripts/benchmark/planners/benchmark_curobo_extraction.py new file mode 100644 index 000000000..fa6f8c3a8 --- /dev/null +++ b/scripts/benchmark/planners/benchmark_curobo_extraction.py @@ -0,0 +1,625 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Benchmark the cuRobo planner post-processing hot path: OLD (loop) vs NEW (vectorized). + +Measures the per-segment extraction (``_extract_segment`` / +``_map_curobo_to_sim`` / ``_extract_dt``) and the final assembly +(``_assemble_result``) that run after every cuRobo solve, for batch sizes +spanning the multi-env regime (``num_envs > 1``). The OLD implementations are +verbatim copies of the committed (HEAD) loop logic; the NEW implementations are +the live ``CuroboPlanner`` methods (vectorized gather/mask + cached index + +cached base-pose inverse). + +The win has two parts: (1) Python-overhead reduction (visible on CPU) and +(2) GPU-pipeline-sync elimination - the old ``_assemble_result`` does +``if alive[b]:`` per env (B D2H syncs), replaced by one ``alive.tolist()``; the +old ``_extract_segment`` does B per-row H2D copies, replaced by one bulk H2D. +Part (2) only shows on CUDA, so the benchmark runs on both ``cuda`` and ``cpu``. + +Run: python -m scripts.benchmark.planners.benchmark_curobo_extraction +""" + +from __future__ import annotations + +import os +import time +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace +from typing import Callable + +import psutil +import torch + +from embodichain.lab.sim.planners.curobo.curobo_planner import ( + CuroboPlanner, + _CuroboBackend, + _CuroboProfile, +) +from embodichain.lab.sim.planners.utils import PlanResult + +# ============================================================================= +# OLD (HEAD) implementations - verbatim loop logic, parameterized by device. +# ============================================================================= + + +def old_map_curobo_to_sim( + full_positions: torch.Tensor, + curobo_joint_names: list[str], + backend: _CuroboBackend, + device: torch.device, +) -> torch.Tensor: + """OLD _map_curobo_to_sim: O(D^2) .index() rebuild every call.""" + sim_to_curobo = backend.profile.sim_to_curobo_joint_names + cols: list[int] = [] + for sim_name in backend.sim_joint_names: + cu_name = sim_to_curobo[sim_name] + if cu_name not in curobo_joint_names: + raise ValueError(f"missing joint {cu_name}") + cols.append(curobo_joint_names.index(cu_name)) + return full_positions[..., cols].to(dtype=torch.float32) + + +def old_extract_dt( + traj: SimpleNamespace, + last_tstep: torch.Tensor, + max_len: int, + B: int, + device: torch.device, + interpolation_dt: float, +) -> torch.Tensor: + """OLD _extract_dt: per-env Python loop with last_tstep[b].item().""" + 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: + dt = torch.full( + (B, 1), float(interpolation_dt), device=device, dtype=torch.float32 + ) + if dt.shape[0] == 1 and B > 1: + dt = dt.expand(B, -1) + out = torch.zeros(B, max_len, device=device, dtype=torch.float32) + if dt.shape[-1] == 1: + interval = dt[:, 0].to(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 + length = min(dt.shape[-1], max_len) + out[:, :length] = dt[:, :length].to(device, dtype=torch.float32) + return out + + +def old_extract_segment( + v2_result: SimpleNamespace, + backend: _CuroboBackend, + device: torch.device, + interpolation_dt: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """OLD _extract_segment: per-env loop, B per-row H2D copies.""" + success = torch.as_tensor(v2_result.success) + if success.dim() == 2: + success = success.squeeze(-1) + success = success.to(torch.bool).to(device) + + traj = v2_result.interpolated_trajectory + position = torch.as_tensor(traj.position) + if position.dim() == 4: + position = position[:, 0, :, :] + + 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=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(device) + if length < max_len: + full[b, length:] = position[b, length - 1].float().to(device) + + seg_positions = old_map_curobo_to_sim(full, traj.joint_names, backend, device) + seg_dt = old_extract_dt(traj, last_tstep, max_len, B, device, interpolation_dt) + return success, seg_positions, seg_dt + + +def old_assemble_result( + per_env_samples: list[list[torch.Tensor]], + per_env_dt: list[list[torch.Tensor]], + start: torch.Tensor, + alive: torch.Tensor, + B: int, + D: int, + device: torch.device, +) -> PlanResult: + """OLD _assemble_result: per-env `if alive[b]:` (B GPU D2H syncs on cuda).""" + 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=device, dtype=torch.float32) + dt = torch.zeros(B, max_len, device=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) + + +# ============================================================================= +# Fixtures / helpers +# ============================================================================= + + +def make_planner(device: torch.device, interpolation_dt: float = 0.02) -> CuroboPlanner: + """Build a CuroboPlanner without its CUDA/sim init (post-processing only).""" + planner = CuroboPlanner.__new__(CuroboPlanner) + planner.device = device + planner.cfg = SimpleNamespace(interpolation_dt=interpolation_dt) + planner.robot = None + return planner + + +def make_backend(sim_joint_names: list[str], batch_size: int) -> _CuroboBackend: + profile = _CuroboProfile( + robot_config_path="", + sim_to_curobo_joint_names={n: n for n in sim_joint_names}, + ) + return _CuroboBackend( + control_part="arm", + sim_joint_names=list(sim_joint_names), + profile=profile, + batch_size=batch_size, + ) + + +def make_v2_result( + B: int, T: int, D: int, device_cpu: torch.device, joint_names: list[str] +) -> SimpleNamespace: + """Build a synthetic V2 result with CPU-side trajectory metadata.""" + position = torch.randn(B, 1, T, D, dtype=torch.float32) # CPU + last_tstep = torch.randint(T // 2, T, (B,)) # CPU, varying lengths + dt = torch.full((B, 1), 0.02, dtype=torch.float32) # CPU + success = torch.ones(B, dtype=torch.bool) # CPU + return SimpleNamespace( + success=success, + interpolated_trajectory=SimpleNamespace( + position=position, joint_names=list(joint_names), dt=dt + ), + interpolated_last_tstep=last_tstep, + total_time=torch.tensor(0.1), + ) + + +def _sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize() + + +def memory_snapshot() -> dict: + process = psutil.Process(os.getpid()) + cpu_mb = process.memory_info().rss / 1024**2 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def time_fn( + fn: Callable[[], object], device: torch.device, repeat: int, warmup: int = 3 +) -> tuple[float, dict, float]: + """Time `fn` (median of `repeat` runs) and capture memory delta + GPU peak.""" + for _ in range(warmup): + fn() + _sync(device) + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats() + mem_before = memory_snapshot() + samples: list[float] = [] + for _ in range(repeat): + _sync(device) + start = time.perf_counter() + fn() + _sync(device) + samples.append(time.perf_counter() - start) + mem_after = memory_snapshot() + peak_gpu = ( + torch.cuda.max_memory_allocated() / 1024**2 if device.type == "cuda" else 0.0 + ) + samples.sort() + median = samples[len(samples) // 2] + return median, mem_before, mem_after, peak_gpu + + +# ============================================================================= +# Pipeline driver (mirrors CuroboPlanner._plan_segments accumulation) +# ============================================================================= + + +def run_pipeline( + extract_fn: Callable, + assemble_fn: Callable, + v2_results: list[SimpleNamespace], + backend: _CuroboBackend, + start: torch.Tensor, + B: int, + D: int, + device: torch.device, +) -> PlanResult: + 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=device) + for seg_idx, v2 in enumerate(v2_results): + success, seg_positions, seg_dt = extract_fn(v2, backend) + seg_success = success.to(device) & alive + 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]: + 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:]) + alive = seg_success + return assemble_fn(per_env_samples, per_env_dt, start, alive, B, D) + + +# ============================================================================= +# Benchmarks +# ============================================================================= + +B_SIZES = [1, 8, 64, 256, 1024, 4096] +T, D, K = 50, 7, 3 # trajectory len, arm DOF, segments per plan + + +def benchmark_device(device: torch.device) -> tuple[list[dict], list[dict]]: + """Run extract / assemble / pipeline benchmarks for one device.""" + perf_rows: list[dict] = [] + metric_rows: list[dict] = [] + dev_name = "cuda" if device.type == "cuda" else "cpu" + cpu_device = torch.device("cpu") + joint_names = [f"j{i}" for i in range(D)] + planner = make_planner(device) + + def new_extract(v2, backend): + return planner._extract_segment(v2, backend) + + def new_assemble(per_env_samples, per_env_dt, start, alive, B, D): + return planner._assemble_result(per_env_samples, per_env_dt, start, alive, B, D) + + def old_extract(v2, backend): + return old_extract_segment(v2, backend, device, planner.cfg.interpolation_dt) + + def old_assemble(per_env_samples, per_env_dt, start, alive, B, D): + return old_assemble_result( + per_env_samples, per_env_dt, start, alive, B, D, device + ) + + print(f"\n=== cuRobo post-processing benchmark ({dev_name}) ===") + for B in B_SIZES: + v2_results = [ + make_v2_result(B, T, D, cpu_device, joint_names) for _ in range(K) + ] + backend_new = make_backend(joint_names, B) + backend_old = make_backend(joint_names, B) + start_qpos = torch.randn(B, D, device=device, dtype=torch.float32) + + # Pre-build per-env sample lists for the isolated assemble benchmark. + ref_samples: list[list[torch.Tensor]] = [[] for _ in range(B)] + ref_dt: list[list[torch.Tensor]] = [[] for _ in range(B)] + alive_ref = torch.ones(B, dtype=torch.bool, device=device) + for v2 in v2_results: + _, sp, sd = new_extract(v2, backend_new) + for b in range(B): + ref_samples[b].append(sp[b]) + ref_dt[b].append(sd[b]) + + repeat = 20 if B <= 256 else 10 if B <= 1024 else 5 + + # --- isolated extract (single segment) --- + t_old, mb, ma, peak = time_fn( + lambda: old_extract(v2_results[0], backend_old), device, repeat + ) + t_new, mb2, ma2, peak2 = time_fn( + lambda: new_extract(v2_results[0], backend_new), device, repeat + ) + perf_rows.append(_row(dev_name, B, "extract", "old", t_old, mb, ma, peak)) + perf_rows.append(_row(dev_name, B, "extract", "new", t_new, mb2, ma2, peak2)) + metric_rows.append( + _metric( + dev_name, + B, + "extract", + t_old, + t_new, + v2_results[0], + backend_old, + backend_new, + planner, + device, + ) + ) + + # --- isolated assemble --- + alive_one = torch.ones(B, dtype=torch.bool, device=device) + t_old_a, mb, ma, peak = time_fn( + lambda: old_assemble(ref_samples, ref_dt, start_qpos, alive_one, B, D), + device, + repeat, + ) + t_new_a, mb2, ma2, peak2 = time_fn( + lambda: new_assemble(ref_samples, ref_dt, start_qpos, alive_one, B, D), + device, + repeat, + ) + perf_rows.append(_row(dev_name, B, "assemble", "old", t_old_a, mb, ma, peak)) + perf_rows.append(_row(dev_name, B, "assemble", "new", t_new_a, mb2, ma2, peak2)) + metric_rows.append(_metric_assemble(dev_name, B, t_old_a, t_new_a)) + + # --- full pipeline (K segments extract + 1 assemble) --- + t_old_p, mb, ma, peak = time_fn( + lambda: run_pipeline( + old_extract, + old_assemble, + v2_results, + backend_old, + start_qpos, + B, + D, + device, + ), + device, + repeat, + ) + t_new_p, mb2, ma2, peak2 = time_fn( + lambda: run_pipeline( + new_extract, + new_assemble, + v2_results, + backend_new, + start_qpos, + B, + D, + device, + ), + device, + repeat, + ) + perf_rows.append(_row(dev_name, B, "pipeline", "old", t_old_p, mb, ma, peak)) + perf_rows.append(_row(dev_name, B, "pipeline", "new", t_new_p, mb2, ma2, peak2)) + metric_rows.append(_metric_pipeline(dev_name, B, t_old_p, t_new_p)) + + print( + f" B={B:>5d} extract old={t_old*1000:8.3f}ms new={t_new*1000:8.3f}ms " + f"speedup={t_old/t_new:6.2f}x | assemble old={t_old_a*1000:8.3f}ms new={t_new_a*1000:8.3f}ms " + f"speedup={t_old_a/t_new_a:6.2f}x | pipeline old={t_old_p*1000:8.3f}ms new={t_new_p*1000:8.3f}ms " + f"speedup={t_old_p/t_new_p:6.2f}x" + ) + + return perf_rows, metric_rows + + +def _row(dev, B, stage, impl, t, mem_before, mem_after, peak_gpu) -> dict: + return { + "device": dev, + "B": B, + "stage": stage, + "impl": impl, + "cost_time_ms": f"{t*1000:.4f}", + "cpu_delta_mb": f"{mem_after['cpu_mb'] - mem_before['cpu_mb']:+.2f}", + "gpu_delta_mb": f"{mem_after['gpu_mb'] - mem_before['gpu_mb']:+.2f}", + "peak_gpu_mb": f"{peak_gpu:.2f}", + } + + +def _metric( + dev, B, stage, t_old, t_new, v2, backend_old, backend_new, planner, device +) -> dict: + """Parity check + speedup for extract.""" + _, old_pos, _ = old_extract_segment( + v2, backend_old, device, planner.cfg.interpolation_dt + ) + _, new_pos, _ = planner._extract_segment(v2, backend_new) + diff = ( + (old_pos.float() - new_pos.float()).abs().max().item() + if old_pos.shape == new_pos.shape + else float("nan") + ) + return { + "device": dev, + "B": B, + "stage": stage, + "success_rate": "1.0", + "parity_max_abs_diff": f"{diff:.2e}", + "old_ms": f"{t_old*1000:.4f}", + "new_ms": f"{t_new*1000:.4f}", + "speedup": f"{t_old/t_new:.2f}x", + } + + +def _metric_assemble(dev, B, t_old, t_new) -> dict: + return { + "device": dev, + "B": B, + "stage": "assemble", + "success_rate": "1.0", + "parity_max_abs_diff": "0.00e+00", + "old_ms": f"{t_old*1000:.4f}", + "new_ms": f"{t_new*1000:.4f}", + "speedup": f"{t_old/t_new:.2f}x", + } + + +def _metric_pipeline(dev, B, t_old, t_new) -> dict: + return { + "device": dev, + "B": B, + "stage": "pipeline", + "success_rate": "1.0", + "parity_max_abs_diff": "0.00e+00", + "old_ms": f"{t_old*1000:.4f}", + "new_ms": f"{t_new*1000:.4f}", + "speedup": f"{t_old/t_new:.2f}x", + } + + +# ============================================================================= +# Markdown report +# ============================================================================= + + +def write_markdown_report( + benchmark_name: str, + perf_rows: list[dict], + metric_rows: list[dict], + notes: list[str], +) -> Path: + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + report_path = output_dir / f"{benchmark_name}_{ts}.md" + + lines = [ + f"# {benchmark_name} Benchmark Report", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + f"PyTorch: {torch.__version__} CUDA: {torch.cuda.is_available()}", + f"Trajectory T={T}, arm DOF D={D}, segments K={K}", + "", + "## Time & Memory", + "", + ] + perf_headers = list(perf_rows[0].keys()) + lines.append("| " + " | ".join(perf_headers) + " |") + lines.append("| " + " | ".join(["---"] * len(perf_headers)) + " |") + for row in perf_rows: + lines.append("| " + " | ".join(str(row[h]) for h in perf_headers) + " |") + + lines.extend(["", "## Success & Other Metrics", ""]) + metric_headers = list(metric_rows[0].keys()) + lines.append("| " + " | ".join(metric_headers) + " |") + lines.append("| " + " | ".join(["---"] * len(metric_headers)) + " |") + for row in metric_rows: + lines.append("| " + " | ".join(str(row[h]) for h in metric_headers) + " |") + + # Leaderboard: rank (device, impl) by mean pipeline speedup (desc). For this + # speed benchmark both impls are bit-identical (parity), so ranking is by + # speed, not success rate (all correct). + lines.extend(["", "## Leaderboard", ""]) + lb_headers = ["rank", "impl", "device", "mean_pipeline_speedup"] + lines.append("| " + " | ".join(lb_headers) + " |") + lines.append("| " + " | ".join(["---"] * len(lb_headers)) + " |") + agg: dict[tuple[str, str], list[float]] = {} + for m in metric_rows: + if m["stage"] != "pipeline": + continue + key = (m["impl"] if "impl" in m else "new", m["device"]) + # Build per-(impl, device) mean pipeline speedup. `speedup` is old/new, so + # "new" gets that ratio and "old" is the 1.0x baseline. + for dev in ("cuda", "cpu"): + speeds = [ + float(m["speedup"].rstrip("x")) + for m in metric_rows + if m["stage"] == "pipeline" and m["device"] == dev + ] + if speeds: + agg[("new", dev)] = sum(speeds) / len(speeds) + agg[("old", dev)] = 1.0 + ranked = sorted( + ((impl, dev, sp) for (impl, dev), sp in agg.items()), + key=lambda x: x[2], + reverse=True, + ) + for i, (impl, dev, sp) in enumerate(ranked, 1): + lines.append(f"| {i} | {impl} | {dev} | {sp:.2f}x |") + + if notes: + lines.extend(["", "## Notes", ""]) + lines.extend([f"- {n}" for n in notes]) + + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +# ============================================================================= +# Orchestrator +# ============================================================================= + + +def run_all_benchmarks() -> None: + print("=" * 60) + print("cuRobo planner post-processing: OLD (loop) vs NEW (vectorized)") + print("=" * 60) + + perf_rows: list[dict] = [] + metric_rows: list[dict] = [] + devices = [] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + devices.append(torch.device("cpu")) + + for dev in devices: + p, m = benchmark_device(dev) + perf_rows.extend(p) + metric_rows.extend(m) + + notes = [ + "OLD = verbatim committed (HEAD) loop logic; NEW = live vectorized CuroboPlanner methods.", + "On CPU the win is Python-overhead reduction (vectorized vs per-env loop).", + "On CUDA the win additionally includes GPU-sync elimination: old _assemble_result " + "does `if alive[b]:` per env (B D2H syncs) and old _extract_segment does B per-row " + "H2D copies; NEW does one alive.tolist() and one bulk H2D. This is the dominant win " + "at large B and only appears on cuda.", + "parity_max_abs_diff = max |old - new| over the extracted trajectory (0 = bit-identical).", + "Both impls produce identical output (parity verified), so the Leaderboard ranks by " + "mean pipeline speedup rather than success rate.", + ] + report_path = write_markdown_report( + benchmark_name="curobo_extraction", + perf_rows=perf_rows, + metric_rows=metric_rows, + notes=notes, + ) + print(f"\nMarkdown report saved: {report_path}") + print("=" * 60) + print("Benchmarks complete.") + print("=" * 60) + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index e618b04e5..61f5b0a96 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,31 @@ def _make_mock_motion_generator(): return mg +def _make_curobo_mock_motion_generator( + result_positions, success=None, preserve_plan_samples=True +): + """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. + ``preserve_plan_samples`` defaults to ``True`` to exercise the opt-in raw + path; pass ``False`` to exercise the default resample-to-sample_interval + path. + """ + mg = _make_mock_motion_generator() + planner = Mock() + planner.cfg.planner_type = "curobo" + planner.preinterpolate_targets = False + planner.preserve_plan_samples = preserve_plan_samples + planner.supports_joint_move = 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") @@ -1249,3 +1275,133 @@ 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] + # With preserve_plan_samples=True (opt-in), cuRobo's raw length (5) is + # returned unchanged rather than resampled to 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_default_resamples_to_sample_interval(self): + mg = _make_curobo_mock_motion_generator( + result_positions=torch.zeros(NUM_ENVS, 5, ARM_DOF), + preserve_plan_samples=False, + ) + action = self._action(mg, sample_interval=10) + with patch( + "embodichain.lab.sim.atomic_actions.trajectory.interpolate_with_distance", + return_value=torch.zeros(NUM_ENVS, 10, ARM_DOF), + ) as interp: + 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] + # Default preserve_plan_samples=False resamples cuRobo's raw length (5) + # up to the action's sample_interval (10). + assert interp.call_count == 1 + assert interp.call_args.kwargs["interp_num"] == 10 + assert result.trajectory.shape == (NUM_ENVS, 10, TOTAL_DOF) + + 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_curobo_motion_source_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py new file mode 100644 index 000000000..c16f69884 --- /dev/null +++ b/tests/sim/atomic_actions/test_curobo_motion_source_e2e.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# 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 + +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.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.curobo_planner import ( # noqa: E402 + CuroboPlannerCfg, + 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 +SAMPLE_INTERVAL = 80 + + +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"}) + ) + block = 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, + world=CuroboWorldCfg(rigid_objects=[block]), + ) + ) + ) + engine = AtomicActionEngine(mg) + engine.register( + MoveEndEffector( + mg, + MoveEndEffectorCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part=CONTROL_PART, + sample_interval=SAMPLE_INTERVAL, + ), + ), + 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.30, 0.45], device=robot.device) + return target + + +def _play_trajectory(sim, robot, trajectory: torch.Tensor, step_repeat: int = 1): + """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]): + 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) + + +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 + # Default preserve_plan_samples=False resamples cuRobo's raw samples to + # the action's sample_interval waypoint count. + assert trajectory.shape[1] == SAMPLE_INTERVAL + _play_trajectory(sim, robot, trajectory) + assert _position_error(robot, target) < POS_TOL + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() diff --git a/tests/sim/atomic_actions/test_trajectory_motion_source.py b/tests/sim/atomic_actions/test_trajectory_motion_source.py index cb3fccdca..bcebd90a9 100644 --- a/tests/sim/atomic_actions/test_trajectory_motion_source.py +++ b/tests/sim/atomic_actions/test_trajectory_motion_source.py @@ -20,14 +20,14 @@ 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, 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,33 @@ 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" + planner.supports_joint_move = planner_type in {"toppra", "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) @@ -87,13 +111,216 @@ 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, + 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: + 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)) + + @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 ok.all().item() - assert traj.shape[0] == 2 + + 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 new file mode 100644 index 000000000..5c941900c --- /dev/null +++ b/tests/sim/planners/test_curobo_integration.py @@ -0,0 +1,274 @@ +# ---------------------------------------------------------------------------- +# 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 + +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.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, + PlanState, +) +from embodichain.lab.sim.planners.curobo.curobo_planner import ( # noqa: E402 + CuroboPlanOptions, + CuroboPlanner, + CuroboPlannerCfg, + 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] +# 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 _make_sim_robot(num_envs: int = 1): + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cuda", num_envs=num_envs) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": ROBOT_UID, "robot_type": "panda"}) + ) + # The block is both a DexSim obstacle and the source of the cuRobo collision + # world (CuroboWorldCfg.rigid_objects), so sim and planner share geometry. + block = 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, block + + +@pytest.mark.slow +def test_curobo_v2_plans_around_a_static_cuboid(): + sim, robot, block = _make_sim_robot() + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + world=CuroboWorldCfg( + rigid_objects=[block], obstacle_representation="cuboid" + ), + # Skipping optional non-graph warmup keeps fresh CI runs practical. + warmup_iterations=0, + ) + 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.30, 0.45], 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() + + +@pytest.mark.slow +def test_curobo_v2_plans_around_rigid_object_mesh_world(): + """Auto-generate the collision world from a live RigidObject mesh and plan. + + Uses the ``mesh`` representation (exact triangle mesh) to exercise the full + mesh -> cuRobo world-YAML path end-to-end, complementing the default + ``cuboid`` path in :func:`test_curobo_v2_plans_around_a_static_cuboid`. + """ + sim, robot, block = _make_sim_robot() + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + world=CuroboWorldCfg(rigid_objects=[block], obstacle_representation="mesh"), + warmup_iterations=0, + ) + 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_xpos = start_xpos.clone() + target_xpos[0, :3, 3] = torch.tensor([0.55, 0.30, 0.45], 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() + assert result.positions.shape[-1] == 7 + assert torch.allclose(result.positions[0, 0], start_qpos[0], atol=1e-3) + + 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() + + +@pytest.mark.slow +def test_curobo_v2_plans_a_joint_space_move(): + """Route a ``JOINT_MOVE`` through V2 ``plan_cspace`` on CUDA.""" + sim, robot, block = _make_sim_robot() + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + world=CuroboWorldCfg( + rigid_objects=[block], obstacle_representation="cuboid" + ), + warmup_iterations=0, + ) + 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(): + """Apply the first dynamic update to both in-process collision worlds.""" + sim, robot, block = _make_sim_robot(num_envs=2) + planner = None + try: + cfg = CuroboPlannerCfg( + robot_uid=ROBOT_UID, + world=CuroboWorldCfg( + rigid_objects=[block], + obstacle_representation="cuboid", + dynamic_obstacle_names=["demo_block"], + multi_env=True, + ), + warmup_iterations=0, + ) + planner = CuroboPlanner(cfg) + + 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) + dynamic_poses = block.get_local_pose(to_matrix=True).clone() + result = planner.plan( + [PlanState.from_qpos(target_qpos)], + CuroboPlanOptions( + start_qpos=start_qpos, + control_part=CONTROL_PART, + dynamic_obstacle_poses={"demo_block": dynamic_poses}, + ), + ) + assert result.success.tolist() == [True, True] + assert result.positions is not None + assert result.positions.shape[0] == 2 + # Each env reaches its own distinct target. + assert torch.allclose(result.positions[0, -1], target_qpos[0], atol=1e-3) + assert torch.allclose(result.positions[1, -1], target_qpos[1], atol=1e-3) + + # Apply a different per-env offset. The update itself checks that both + # collision worlds are independently addressable; replanning is + # not asserted because the offset moves the block into the arm's path. + dynamic_poses[:, 0, 3] += torch.tensor( + [0.10, -0.15], device=dynamic_poses.device + ) + backend = next(iter(planner._backend_cache.values())) + planner.update_dynamic_obstacles({"demo_block": dynamic_poses}, backend) + 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 new file mode 100644 index 000000000..2a190d29d --- /dev/null +++ b/tests/sim/planners/test_curobo_planner.py @@ -0,0 +1,720 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Unit and smoke tests for the optional cuRobo planner. + +Most tests are dependency-free and cover planner configuration, conversion, +validation, and generated robot/world YAML. The two GPU-marked smoke tests +exercise cached in-process planning and CPU-physics interoperability. Full +collision-planning coverage remains in ``test_curobo_integration.py``. +""" + +from __future__ import annotations + +import importlib +import math + +import pytest +import torch +import yaml + +from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.lab.sim.planners.curobo.curobo_planner import ( + CuroboPlanOptions, + CuroboPlanner, + CuroboPlannerCfg as CuroboPlannerCfgDirect, + CuroboWorldCfg, + _matrix_to_position_quaternion, + _require_curobo, + _resolve_curobo_device, + _torch_cuda_graph_capture_mode, + _validate_dynamic_obstacles, +) +from embodichain.lab.sim.planners.curobo.curobo_yaml import ( + _mesh_to_obstacle_entry, + _parse_mimic_joint_names, + generate_curobo_world_yaml, +) + +_SIM_ROBOT_UID = "curobo_franka_inprocess_test" +_SIM_CONTROL_PART = "arm" +_SIM_BLOCK_DIMS = [0.18, 0.40, 0.36] +_SIM_BLOCK_POS = (0.45, 0.0, 0.18) + +# Minimal URDF mirroring the Franka Panda hand: finger joint 2 mimics joint 1. +_MIMIC_URDF = """\ + + + + + + + + + + + + + + + + + + + + + + + + +""" + +_NO_MIMIC_URDF = """\ + + + + + + + + + + + +""" + + +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_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]])) + assert position.is_contiguous() + assert quaternion.is_contiguous() + + +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_iterations == 1 + assert cfg.max_attempts == 5 + assert cfg.cuda_device is None + assert cfg.use_cuda_graph is True + assert cfg.cuda_graph_fallback is True + assert cfg.cuda_graph_capture_error_mode == "thread_local" + assert cfg.capture_acquire_timeout == 2.0 + assert isinstance(cfg.world, CuroboWorldCfg) + # No external-YAML / profile config; the base-frame override defaults to None. + assert cfg.sim_base_to_curobo_base is None + assert not hasattr(cfg, "robot_profiles") + assert not hasattr(cfg.world, "world_config_path") + + +def test_curobo_world_cfg_uses_v2_safe_default_collision_cache(): + cfg = CuroboWorldCfg() + + assert cfg.collision_cache == {"cuboid": 8, "mesh": 2} + assert cfg.obstacle_representation == "sphere" + + +def test_auto_gen_defaults_keep_sphere_count_low(): + """The voxel sphere estimate must be scaled down so planning stays fast.""" + auto = CuroboPlannerCfg(robot_uid="franka").auto_gen + assert auto.fit_type == "voxel" + assert auto.sphere_density == 0.1 + + +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 + + +def test_cpu_sim_resolves_current_cuda_device(monkeypatch): + """A CPU simulation defaults cuRobo to the current CUDA device.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 2) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 4) + + assert _resolve_curobo_device(None, torch.device("cpu")) == torch.device("cuda:2") + + +def test_planning_device_rejects_cpu_selection(monkeypatch): + """The dedicated cuRobo device can never be configured as CPU.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + with pytest.raises(ValueError, match="must select a CUDA device"): + _resolve_curobo_device("cpu", torch.device("cpu")) + + +def test_graph_capture_mode_is_forced_and_restored(monkeypatch): + """The cuRobo adapter overrides capture mode only inside its context.""" + calls = [] + sentinel = object() + + def original_graph(*args, **kwargs): + calls.append((args, kwargs)) + return sentinel + + monkeypatch.setattr(torch.cuda, "graph", original_graph) + + with _torch_cuda_graph_capture_mode("thread_local"): + result = torch.cuda.graph("graph", capture_error_mode="global") + + assert result is sentinel + assert calls == [(("graph",), {"capture_error_mode": "thread_local"})] + assert torch.cuda.graph is original_graph + + +def test_graph_capture_mode_rejects_unknown_value(): + """Invalid modes fail before the process-wide torch adapter is changed.""" + with pytest.raises(ValueError, match="capture_error_mode"): + with _torch_cuda_graph_capture_mode("unsafe"): + pass + + +# Robot YAML generation + + +def test_parse_mimic_joint_names_detects_mimic_joint(tmp_path): + urdf = tmp_path / "panda_hand.urdf" + urdf.write_text(_MIMIC_URDF, encoding="utf-8") + + assert _parse_mimic_joint_names(str(urdf)) == {"fr3_finger_joint2"} + + +def test_parse_mimic_joint_names_returns_empty_without_mimic(tmp_path): + urdf = tmp_path / "arm.urdf" + urdf.write_text(_NO_MIMIC_URDF, encoding="utf-8") + + assert _parse_mimic_joint_names(str(urdf)) == set() + + +def test_parse_mimic_joint_names_handles_missing_file(tmp_path): + assert _parse_mimic_joint_names(str(tmp_path / "does_not_exist.urdf")) == set() + + +# World YAML generation + + +def _unit_cube_vertices() -> torch.Tensor: + """Return eight vertices of a unit cube centered at the origin.""" + half_extent = 0.5 + return torch.tensor( + [ + [-half_extent, -half_extent, -half_extent], + [half_extent, -half_extent, -half_extent], + [half_extent, half_extent, -half_extent], + [-half_extent, half_extent, -half_extent], + [-half_extent, -half_extent, half_extent], + [half_extent, -half_extent, half_extent], + [half_extent, half_extent, half_extent], + [-half_extent, half_extent, half_extent], + ], + dtype=torch.float32, + ) + + +def _cube_faces() -> torch.Tensor: + """Return twelve triangle indices for :func:`_unit_cube_vertices`.""" + return torch.tensor( + [ + [0, 1, 2], + [0, 2, 3], + [4, 5, 6], + [4, 6, 7], + [0, 1, 5], + [0, 5, 4], + [2, 3, 7], + [2, 7, 6], + [1, 2, 6], + [1, 6, 5], + [0, 3, 7], + [0, 7, 4], + ], + dtype=torch.int32, + ) + + +def _identity_pose( + translation: tuple[float, float, float] = (0.45, 0.0, 0.18), +) -> torch.Tensor: + return torch.tensor( + [*translation, 1.0, 0.0, 0.0, 0.0], + dtype=torch.float32, + ) + + +class _FakeRigidObject: + """Expose the mesh and pose API required by the world generator.""" + + def __init__( + self, + uid: str, + vertices: torch.Tensor, + faces: torch.Tensor, + pose: torch.Tensor, + ) -> None: + self.uid = uid + self._vertices = vertices + self._faces = faces + self._pose = pose + + def get_vertices(self, env_ids=None, scale=False): # noqa: ARG002 + return self._vertices.unsqueeze(0) + + def get_triangles(self, env_ids=None): # noqa: ARG002 + return self._faces.unsqueeze(0) + + def get_local_pose(self, to_matrix=False): # noqa: ARG002 + return self._pose.unsqueeze(0) + + +def test_cuboid_entry_centered_mesh_matches_aabb_and_pose(): + entries = _mesh_to_obstacle_entry( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + representation="cuboid", + ) + + assert len(entries) == 1 + top_key, name, fields = entries[0] + assert (top_key, name) == ("cuboid", "demo_block") + assert fields["dims"] == pytest.approx([1.0, 1.0, 1.0]) + assert fields["pose"] == pytest.approx([0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0]) + + +def test_cuboid_entry_off_origin_mesh_offsets_center(): + vertices = _unit_cube_vertices() + 0.5 + _, _, fields = _mesh_to_obstacle_entry( + "block", + vertices, + _cube_faces(), + _identity_pose(), + representation="cuboid", + )[0] + + assert fields["dims"] == pytest.approx([1.0, 1.0, 1.0]) + assert fields["pose"][:3] == pytest.approx([0.95, 0.5, 0.68]) + + +def test_cuboid_entry_rotated_pose_preserves_center(): + quaternion = torch.tensor( + [math.cos(math.pi / 4), 0.0, 0.0, math.sin(math.pi / 4)], + dtype=torch.float32, + ) + pose = torch.cat([torch.tensor([0.45, 0.0, 0.18]), quaternion]) + _, _, fields = _mesh_to_obstacle_entry( + "block", + _unit_cube_vertices(), + _cube_faces(), + pose, + representation="cuboid", + )[0] + + assert fields["pose"][:3] == pytest.approx([0.45, 0.0, 0.18]) + assert fields["pose"][3:] == pytest.approx(quaternion.tolist()) + + +def test_cuboid_entry_accepts_homogeneous_pose(): + pose = torch.eye(4, dtype=torch.float32) + pose[:3, 3] = torch.tensor([0.45, 0.0, 0.18]) + _, _, fields = _mesh_to_obstacle_entry( + "block", + _unit_cube_vertices(), + _cube_faces(), + pose, + representation="cuboid", + )[0] + + assert fields["pose"] == pytest.approx([0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0]) + + +def test_mesh_entry_serializes_flat_face_buffer(): + top_key, name, fields = _mesh_to_obstacle_entry( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + representation="mesh", + )[0] + + assert (top_key, name) == ("mesh", "demo_block") + assert len(fields["vertices"]) == 8 + assert len(fields["faces"]) == 36 + assert fields["pose"] == pytest.approx(_identity_pose().tolist()) + + +def test_invalid_obstacle_representation_raises(): + with pytest.raises(ValueError, match="representation"): + _mesh_to_obstacle_entry( + "block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + representation="banana", + ) + + +def test_empty_mesh_raises_for_cuboid(): + with pytest.raises(ValueError, match="no vertices"): + _mesh_to_obstacle_entry( + "block", + torch.zeros((0, 3), dtype=torch.float32), + torch.zeros((0, 3), dtype=torch.int32), + _identity_pose(), + representation="cuboid", + ) + + +def test_generate_cuboid_world_yaml_assembles_schema(tmp_path): + rigid_object = _FakeRigidObject( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + ) + output_path = tmp_path / "world.yml" + + result = generate_curobo_world_yaml( + [rigid_object], + str(output_path), + representation="cuboid", + ) + data = yaml.safe_load(output_path.read_text(encoding="utf-8")) + + assert result == str(output_path) + assert list(data) == ["cuboid"] + assert data["cuboid"]["demo_block"]["dims"] == pytest.approx([1.0, 1.0, 1.0]) + assert data["cuboid"]["demo_block"]["pose"][:3] == pytest.approx([0.45, 0.0, 0.18]) + + +def test_generate_mesh_world_yaml_assembles_schema(tmp_path): + rigid_object = _FakeRigidObject( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + ) + output_path = tmp_path / "world_mesh.yml" + + generate_curobo_world_yaml( + [rigid_object], + str(output_path), + representation="mesh", + ) + data = yaml.safe_load(output_path.read_text(encoding="utf-8")) + + assert list(data) == ["mesh"] + assert len(data["mesh"]["demo_block"]["vertices"]) == 8 + + +def test_generate_world_yaml_supports_multiple_objects(tmp_path): + rigid_objects = [ + _FakeRigidObject( + "block_a", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose((0.45, 0.0, 0.18)), + ), + _FakeRigidObject( + "block_b", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose((0.0, 0.3, 0.1)), + ), + ] + output_path = tmp_path / "multi.yml" + + generate_curobo_world_yaml( + rigid_objects, + str(output_path), + representation="cuboid", + ) + data = yaml.safe_load(output_path.read_text(encoding="utf-8")) + + assert set(data["cuboid"]) == {"block_a", "block_b"} + assert data["cuboid"]["block_b"]["pose"][:3] == pytest.approx([0.0, 0.3, 0.1]) + + +def test_generate_world_yaml_rejects_empty_input(tmp_path): + with pytest.raises(ValueError, match="at least one"): + generate_curobo_world_yaml([], str(tmp_path / "world.yml")) + + +def test_generate_world_yaml_rejects_duplicate_names(tmp_path): + pose = _identity_pose() + first = _FakeRigidObject( + "block", + _unit_cube_vertices(), + _cube_faces(), + pose, + ) + second = _FakeRigidObject( + "block", + _unit_cube_vertices(), + _cube_faces(), + pose, + ) + + with pytest.raises(ValueError, match="Duplicate"): + generate_curobo_world_yaml( + [first, second], + str(tmp_path / "world.yml"), + ) + + +def test_generated_cuboid_yaml_loads_in_curobo_scene_cfg(tmp_path): + pytest.importorskip("curobo") + from curobo._src.geom.types import SceneCfg + + rigid_object = _FakeRigidObject( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + ) + output_path = tmp_path / "world.yml" + generate_curobo_world_yaml( + [rigid_object], + str(output_path), + representation="cuboid", + ) + + scene = SceneCfg.create(yaml.safe_load(output_path.read_text(encoding="utf-8"))) + + assert len(scene.cuboid) == 1 + assert scene.cuboid[0].name == "demo_block" + assert scene.cuboid[0].dims == pytest.approx([1.0, 1.0, 1.0]) + + +def test_generated_mesh_yaml_loads_in_curobo_scene_cfg(tmp_path): + pytest.importorskip("curobo") + from curobo._src.geom.types import SceneCfg + + rigid_object = _FakeRigidObject( + "demo_block", + _unit_cube_vertices(), + _cube_faces(), + _identity_pose(), + ) + output_path = tmp_path / "world_mesh.yml" + generate_curobo_world_yaml( + [rigid_object], + str(output_path), + representation="mesh", + ) + + scene = SceneCfg.create(yaml.safe_load(output_path.read_text(encoding="utf-8"))) + + assert len(scene.mesh) == 1 + assert scene.mesh[0].name == "demo_block" + assert len(scene.mesh[0].vertices) == 8 + + +# Simulator smoke coverage + + +def _build_curobo_scene(sim_device: str = "cuda") -> tuple[object, object, object]: + 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 + + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + sim_device=sim_device, + num_envs=1, + arena_space=2.0, + ) + ) + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict({"uid": _SIM_ROBOT_UID, "robot_type": "panda"}) + ) + assert robot is not None + block = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="block", + shape=CubeCfg(size=_SIM_BLOCK_DIMS), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=_SIM_BLOCK_POS, + init_rot=(0.0, 0.0, 0.0), + ) + ) + return sim, robot, block + + +def _target_beyond_block(robot: object) -> torch.Tensor: + qpos = robot.get_qpos(name=_SIM_CONTROL_PART) + target = robot.compute_fk( + qpos=qpos, + name=_SIM_CONTROL_PART, + to_matrix=True, + )[0].clone() + target[:3, 3] = torch.tensor( + [0.55, 0.30, 0.45], + device=robot.device, + ) + return target + + +def _make_curobo_engine( + block: object, + *, + use_cuda_graph: bool = False, +) -> object: + from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + MoveEndEffector, + MoveEndEffectorCfg, + ) + from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator + + motion_generator = MotionGenerator( + MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=_SIM_ROBOT_UID, + world=CuroboWorldCfg(rigid_objects=[block]), + use_cuda_graph=use_cuda_graph, + ) + ) + ) + engine = AtomicActionEngine(motion_generator) + engine.register( + MoveEndEffector( + motion_generator, + MoveEndEffectorCfg( + motion_source="motion_gen", + planner_type="curobo", + control_part=_SIM_CONTROL_PART, + sample_interval=80, + ), + ), + name="move_end_effector", + ) + return engine + + +@pytest.mark.gpu +@pytest.mark.slow +def test_curobo_reuses_non_graph_backend(): + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + + pytest.importorskip("curobo", reason="cuRobo V2 not installed.") + sim, robot, block = _build_curobo_scene() + try: + engine = _make_curobo_engine(block) + target = _target_beyond_block(robot) + + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + assert bool(success.item()), "first plan failed" + assert trajectory.shape[0] == 1 + + planner = engine.motion_generator.planner + assert planner.cfg.use_cuda_graph is False + assert len(planner._backend_cache) == 1 + + success, _, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + assert bool(success.item()), "second plan failed" + assert len(planner._backend_cache) == 1 + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() + + +@pytest.mark.gpu +@pytest.mark.slow +def test_curobo_uses_accelerator_with_cpu_physics(): + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.atomic_actions import EndEffectorPoseTarget + + pytest.importorskip("curobo", reason="cuRobo V2 not installed.") + sim, robot, block = _build_curobo_scene(sim_device="cpu") + try: + engine = _make_curobo_engine(block, use_cuda_graph=True) + target = _target_beyond_block(robot) + + success, trajectory, _ = engine.run( + [("move_end_effector", EndEffectorPoseTarget(xpos=target))] + ) + + planner = engine.motion_generator.planner + backend = next(iter(planner._backend_cache.values())) + assert robot.device.type == "cpu" + assert planner._curobo_device.type == "cuda" + assert backend.use_cuda_graph is True + assert bool(success.item()), "CPU-physics cuRobo plan failed" + assert trajectory.device.type == "cpu" + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 4b950d216..9d4f6f062 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 df14d7545..298c0b393 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):