diff --git a/docs/pages/advanced/motion_planners.rst b/docs/pages/advanced/motion_planners.rst index 84f9460..7080dd2 100644 --- a/docs/pages/advanced/motion_planners.rst +++ b/docs/pages/advanced/motion_planners.rst @@ -2,9 +2,34 @@ Motion Planners =============== SkillGen plans collision-free transit motions with a pluggable motion-planner backend. The -shipped backend is `cuRobo `_ — GPU-accelerated, collision-aware +shipped backends both wrap `cuRobo `_ — GPU-accelerated, collision-aware trajectory optimization. Planners live in ``isaac_autodata_interfaces/motion_planners/``. +.. list-table:: + :widths: 20 20 60 + :header-rows: 1 + + * - ``--planner_backend`` + - Classes + - cuRobo version + * - ``curobo`` (default) + - ``CuroboPlanner`` / ``CuroboPlannerCfg`` + - v1 (``curobo.wrap.reacher.motion_gen``) + * - ``curobo_v2`` + - ``CuroboV2Planner`` / ``CuroboV2PlannerCfg`` + - v2 (``curobo.motion_planner.MotionPlanner``) + +Both cuRobo versions install as the same ``curobo`` package at incompatible versions, so each +backend needs its own environment and only the selected one is ever imported. Resolve a backend +by name rather than importing a class directly: + +.. code-block:: python + + from isaac_autodata_interfaces.motion_planners import get_planner_backend + + planner_cls, config_cls = get_planner_backend("curobo_v2") + planner = planner_cls(datastream=datastream, config=config_cls.from_profile("franka_stack_cube"), env_id=0) + The Planner Interface --------------------- @@ -36,19 +61,23 @@ constructed per environment (``--num_envs`` planners total). cuRobo Backend -------------- -``CuroboPlanner`` is configured by ``CuroboPlannerCfg``. Configurations are resolved from -the task id: +``CuroboPlanner`` is configured by ``CuroboPlannerCfg`` (and ``CuroboV2Planner`` by +``CuroboV2PlannerCfg``). Both configs expose the same two selectors, so switching backends +changes no other call site: .. code-block:: python from isaac_autodata_interfaces.motion_planners.curobo.curobo_planner_cfg import CuroboPlannerCfg config = CuroboPlannerCfg.from_task_name("Isaac-Stack-Cube-Franka-IK-Rel-v0") + config = CuroboPlannerCfg.from_profile("franka_stack_cube_bin") ``from_task_name()`` pattern-matches the task id to a named preset (e.g. ``franka_config()``, ``franka_stack_cube_bin_config()``); unknown robots fall back to the -Franka preset with a printed warning. The ``generate_dataset.py`` entry point calls this -automatically when ``--alg skillgen`` is selected — no CLI flags needed. +Franka preset. ``from_profile()`` takes a profile name directly, which is how an +:doc:`environment profile <../concepts/environment_profiles>` pins planner tuning to its +scene; both backends register the same profile names. The ``generate_dataset.py`` entry +point resolves the config automatically when ``--alg skillgen`` is selected. Key configuration fields: @@ -94,10 +123,13 @@ Visualization and Debugging --------------------------- The cuRobo backend can visualize its collision-sphere model and planned trajectories via -`rerun `_ (``visualize_spheres`` / ``visualize_plan`` config flags). -During multi-env generation these are enabled only for env 0 to keep the simulation -responsive. Planner diagnostics (success rates, timing) are available through -``get_planner_info()``. +`rerun `_ (``visualize_spheres`` / ``visualize_plan`` config flags, or +``--visualize_plan`` on the CLI). During multi-env generation these are enabled only for env 0 +to keep the simulation responsive. The v2 backend renders in the robot base frame — the frame +cuRobo collision-checks in — so spheres, obstacles, and the goal marker line up; it draws +obstacles as their real meshes and colors attached-object spheres separately from the robot's +own. ``visualize_spheres`` (in-sim sphere spawning) is v1-only. Planner diagnostics are +available through ``get_planner_info()``. Adding a New Backend -------------------- @@ -107,5 +139,7 @@ Adding a New Backend 2. Construct your planners where SkillGen expects them — one per env id, exposing ``update_world_and_plan_motion(...)`` and ``get_planned_poses()`` (this is all ``SkillGen`` requires of a planner). -3. Wire construction into your entry point the way ``generate_dataset.py`` builds its cuRobo - planners (``_build_motion_planners``). +3. Register the backend in ``motion_planners/__init__.py`` (``_BACKEND_SPECS``) so + ``get_planner_backend()`` resolves it by name, and add the name to + ``generate_dataset.py``'s ``--planner_backend`` choices. Import your planner lazily from + the subpackage's ``__init__``, so selecting a different backend never imports yours. diff --git a/docs/pages/workflows/skillgen/index.rst b/docs/pages/workflows/skillgen/index.rst index a3a7f11..6b6629e 100644 --- a/docs/pages/workflows/skillgen/index.rst +++ b/docs/pages/workflows/skillgen/index.rst @@ -246,6 +246,11 @@ When motion planning fails for a trial — no collision-free path to the skill s trial is abandoned and counted as a failure; with ``guarantee_success: true`` generation simply retries with a new scene configuration until the trial target is met. +``--planner_backend`` selects which cuRobo version plans those motions: ``curobo`` (v1, +the default) or ``curobo_v2``. Both accept every flag shown above; each needs its matching +cuRobo version installed, so they live in separate environments. See +:doc:`../../advanced/motion_planners`. + For a full-scale run, raise ``--generation_num_trials`` (hundreds to thousands for policy training) and keep ``--viz none`` — rendering slows generation considerably. See `Performance and Scaling`_ before choosing ``--num_envs``. diff --git a/isaac_autodata_interfaces/motion_planners/__init__.py b/isaac_autodata_interfaces/motion_planners/__init__.py index 6eb5fc7..82e2331 100644 --- a/isaac_autodata_interfaces/motion_planners/__init__.py +++ b/isaac_autodata_interfaces/motion_planners/__init__.py @@ -3,34 +3,78 @@ """Motion-planner backends for SkillGen. -This branch ships the cuRobo v1 backend only. The package exposes the abstract -:class:`MotionPlannerBase` and the v1 :class:`CuroboPlanner` / :class:`CuroboPlannerCfg`. +Two cuRobo backends ship here, selected by name through :func:`get_planner_backend`: -The planner class is imported lazily so the configuration dataclass and the abstract base can be -loaded in sim-free contexts (e.g. unit tests or CLI argument parsing) without pulling in cuRobo -or Isaac Lab. +* ``"curobo"`` — :class:`CuroboPlanner` / :class:`CuroboPlannerCfg`, built on cuRobo v1. +* ``"curobo_v2"`` — :class:`CuroboV2Planner` / :class:`CuroboV2PlannerCfg`, built on cuRobo v2. + +Both implement :class:`MotionPlannerBase` and accept the same profile names, so an entry point +picks a backend without changing any other logic. + +Nothing is imported at module load. Both cuRobo versions install under the ``curobo`` package +name at incompatible versions and live in separate environments, so importing a backend that is +not installed would fail. Resolution waits until a backend is requested. """ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from isaac_autodata_interfaces.motion_planners.curobo import CuroboPlannerCfg from isaac_autodata_interfaces.motion_planners.motion_planner_base import MotionPlannerBase if TYPE_CHECKING: - from isaac_autodata_interfaces.motion_planners.curobo import CuroboPlanner + from isaac_autodata_interfaces.motion_planners.curobo import CuroboPlanner, CuroboPlannerCfg + from isaac_autodata_interfaces.motion_planners.curobo_v2 import CuroboV2Planner, CuroboV2PlannerCfg __all__ = [ "MotionPlannerBase", "CuroboPlanner", "CuroboPlannerCfg", + "CuroboV2Planner", + "CuroboV2PlannerCfg", + "PLANNER_BACKENDS", + "get_planner_backend", ] +# Backend name -> (subpackage, planner class name, config class name). +_BACKEND_SPECS: dict[str, tuple[str, str, str]] = { + "curobo": ("curobo", "CuroboPlanner", "CuroboPlannerCfg"), + "curobo_v2": ("curobo_v2", "CuroboV2Planner", "CuroboV2PlannerCfg"), +} + +PLANNER_BACKENDS: tuple[str, ...] = tuple(_BACKEND_SPECS) +"""Names of the selectable planner backends, in registration order. + +``"curobo"`` is the cuRobo v1 backend and the default; ``"curobo_v2"`` is the cuRobo v2 backend. +""" + + +def get_planner_backend(name: str) -> tuple[type[MotionPlannerBase], type]: + """Resolve a backend name to its planner and configuration classes. + + Importing the backend pulls in the matching cuRobo version, so only the requested backend + is loaded. + + Args: + name: Backend name; one of :data:`PLANNER_BACKENDS`. + + Returns: + The ``(planner_class, config_class)`` pair for the backend. The planner takes + ``(datastream, config, env_id)`` and the config exposes ``from_profile`` / + ``from_task_name``. + """ + assert name in _BACKEND_SPECS, f"Unknown planner backend {name!r}. Available: {list(PLANNER_BACKENDS)}" + import importlib + + subpackage, planner_name, config_name = _BACKEND_SPECS[name] + module = importlib.import_module(f"{__name__}.{subpackage}") + return getattr(module, planner_name), getattr(module, config_name) + -def __getattr__(name: str): - if name == "CuroboPlanner": - from isaac_autodata_interfaces.motion_planners.curobo import CuroboPlanner as _CuroboPlanner +def __getattr__(name: str) -> Any: + for subpackage, planner_name, config_name in _BACKEND_SPECS.values(): + if name in (planner_name, config_name): + import importlib - return _CuroboPlanner - raise AttributeError(f"module 'motion_planners' has no attribute {name!r}") + return getattr(importlib.import_module(f"{__name__}.{subpackage}"), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/isaac_autodata_interfaces/motion_planners/curobo/plan_visualizer.py b/isaac_autodata_interfaces/motion_planners/curobo/plan_visualizer.py index 7183344..4f131c4 100644 --- a/isaac_autodata_interfaces/motion_planners/curobo/plan_visualizer.py +++ b/isaac_autodata_interfaces/motion_planners/curobo/plan_visualizer.py @@ -405,7 +405,6 @@ def _clear_visualization(self) -> None: for entity in entities: rr.log(f"world/{entity_type}/{entity}", rr.Clear(recursive=True)) self._sphere_entities[entity_type] = [] - self._current_frame = 0 def clear_visualization(self) -> None: """Public method to clear the visualization.""" diff --git a/isaac_autodata_interfaces/motion_planners/curobo_v2/__init__.py b/isaac_autodata_interfaces/motion_planners/curobo_v2/__init__.py new file mode 100644 index 0000000..faa2301 --- /dev/null +++ b/isaac_autodata_interfaces/motion_planners/curobo_v2/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""cuRobo v2 motion-planner backend. + +Implements :class:`MotionPlannerBase` on :class:`curobo.motion_planner.MotionPlanner`. + +The planner is imported lazily so the configuration can be loaded without a simulator, for +backend selection or CLI argument parsing. Both cuRobo versions install under the ``curobo`` +package name, so only the selected backend is ever imported. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaac_autodata_interfaces.motion_planners.curobo_v2.curobo_v2_planner import CuroboV2Planner + +from isaac_autodata_interfaces.motion_planners.curobo_v2.curobo_v2_planner_cfg import CuroboV2PlannerCfg + +__all__ = [ + "CuroboV2Planner", + "CuroboV2PlannerCfg", +] + + +def __getattr__(name: str): + if name == "CuroboV2Planner": + from isaac_autodata_interfaces.motion_planners.curobo_v2.curobo_v2_planner import ( + CuroboV2Planner as _CuroboV2Planner, + ) + + return _CuroboV2Planner + raise AttributeError(f"module 'curobo_v2' has no attribute {name!r}") diff --git a/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner.py b/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner.py new file mode 100644 index 0000000..991913e --- /dev/null +++ b/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner.py @@ -0,0 +1,1116 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Motion planner backed by cuRobo v2. + +Implements :class:`MotionPlannerBase` on :class:`curobo.motion_planner.MotionPlanner`, letting +:mod:`isaac_autodata_core.algorithms.SkillGen` plan collision-free transits between skill +segments. + +Each planner instance serves one environment and covers: + +* Collision geometry extracted once from the USD stage and kept in sync per plan. +* Attaching a grasped object to the robot so it is collision-checked as part of the arm. +* Planning to a target end-effector pose in three phases: retreat, approach, goal. + +All world state is read through the :class:`Datastream` facade, so the planner stays +independent of the simulator wiring. + +Two conventions are translated at the cuRobo boundary. Callers exchange poses in the +env-relative frame while cuRobo plans in the robot base frame, and Isaac Lab stores +quaternions as ``(x, y, z, w)`` while cuRobo stores ``(w, x, y, z)``. +""" + +from __future__ import annotations + +import contextlib +import logging +import numpy as np +import torch +from typing import TYPE_CHECKING, Any + +# The sphere-fit enum and USD scene parser have no public alias, but the public +# ``curobo.motion_planner.MotionPlanner`` API consumes both directly. +from curobo._src.geom.sphere_fit.types import SphereFitType +from curobo._src.util.usd_scene_parser import UsdSceneParser +from curobo.motion_planner import MotionPlanner, MotionPlannerCfg +from curobo.scene import Cuboid, Scene +from curobo.types import GoalToolPose, JointState, Pose + +from isaac_autodata_interfaces.motion_planners.curobo_v2.curobo_v2_planner_cfg import CuroboV2PlannerCfg +from isaac_autodata_interfaces.motion_planners.motion_planner_base import MotionPlannerBase + +if TYPE_CHECKING: + from isaac_autodata_interfaces.datastream.datastream import Datastream + + +class CuroboV2Planner(MotionPlannerBase): + """cuRobo v2 backend implementing :class:`MotionPlannerBase`. + + Args: + datastream: Read facade over the env, task, and embodiment. All world state (collision + geometry source, object poses, joint configuration) is read through it. + config: Backend-specific configuration; see :class:`CuroboV2PlannerCfg`. + env_id: Index of the env this planner instance serves in a vectorized setup. + debug: Whether to print detailed debug information during planning. + """ + + # Each planner owns a single-world MotionPlanner (multi_env=False), so every collision-world + # operation targets slot 0. ``env_id`` selects which Isaac env to read state from. + _COLLISION_ENV_IDX = 0 + + def __init__( + self, + datastream: Datastream, + config: CuroboV2PlannerCfg, + env_id: int = 0, + debug: bool = False, + ) -> None: + super().__init__(datastream=datastream, env_id=env_id, debug=debug) + self.config = config + # Tag the logger with the env id: planners for every env run concurrently and their + # messages interleave, so an untagged line cannot be traced back to one. + self._logger = logging.getLogger(f"CuroboV2Planner_{env_id}") + + self._logger.info("Building cuRobo MotionPlanner for %r", config.robot_name) + planner_kwargs = config.to_v2_kwargs() + planner_kwargs["robot"] = self._load_robot_config_dict() + self.motion_planner = MotionPlanner(MotionPlannerCfg.create(**planner_kwargs)) + self.motion_planner.warmup() + self._logger.info("MotionPlanner ready. tool_frames=%s", self.motion_planner.tool_frames) + + if self.config.visualize_spheres: + self._logger.warning("visualize_spheres is not implemented; use visualize_plan instead.") + + # Gripper kinematics for the open and closed poses, plus the pose currently loaded into + # the planner. See :meth:`_set_gripper_state`. + self._gripper_kinematics: dict[bool, Any] | None = self._build_gripper_kinematics() + self._gripper_is_closed: bool | None = None + + # Current plan. The end-effector poses are computed once when the plan is stored so + # waypoint iteration does not re-run forward kinematics for every waypoint. + self._planned_joint_trajectory: JointState | None = None + self._planned_eef_poses: list[torch.Tensor] = [] + self._waypoint_index: int = 0 + self._currently_attached: str | None = None + + # Transforms between the env-relative frame callers use and the robot base frame cuRobo + # plans in, refreshed each plan by :meth:`_refresh_base_frame_transform`. + self._base_pose_in_env: torch.Tensor | None = None + self._env_to_base_transform: torch.Tensor | None = None + self._base_frame_is_identity: bool = True + + # Joints the solvers optimize over. A robot config may lock joints (e.g. Franka's + # fingers) that still appear in ``MotionPlanner.joint_names`` but are excluded from the + # optimization, so states passed to :meth:`plan_pose` must omit them. + self._active_joint_names: list[str] = self._compute_active_joint_names() + self._logger.info( + "Active DoF: %d / %d %s", + len(self._active_joint_names), + len(self.motion_planner.joint_names), + self._active_joint_names, + ) + + # Collision world, built once from the USD stage. ``_object_mapping`` relates scene + # object names to cuRobo obstacle names so per-plan pose updates know what to move; + # objects matching ``_static_object_substrings`` are treated as fixed and never moved. + self._object_mapping: dict[str, str] = {} + self._world_obstacle_names: list[str] = [] + self._static_object_substrings: list[str] = [s.lower() for s in self.config.static_objects] + self._obstacle_storage_cache: dict[str, Any] = {} + self._attached_curobo_name: str | None = None + self._initialize_static_world() + + # Imported lazily so runs without visualization do not require the rerun-sdk dependency. + self.plan_visualizer: Any = None + if self.config.visualize_plan: + from isaac_autodata_interfaces.motion_planners.curobo_v2.plan_visualizer import PlanVisualizer + + self.plan_visualizer = PlanVisualizer( + robot_name=self.config.robot_name or "robot", + recording_id=f"curobo_v2_plan_{self.env_id}", + save_path=f"curobo_v2_plan_env{self.env_id}.rrd", + debug=self.debug, + ) + self.plan_visualizer.set_motion_planner_reference(self.motion_planner) + + def _load_robot_config_dict(self, lock_joints: dict[str, float] | None = None) -> dict[str, Any]: + """Load the robot config, applying the sphere allocation and locked-joint overrides. + + cuRobo sizes each link's collision-sphere buffer when the config loads, and an attach + asking for more spheres than the link was allocated fails. Since + :meth:`MotionPlannerCfg.create` accepts a config dict as readily as a path, the overrides + are applied here instead of maintaining a separate robot YAML. + + Args: + lock_joints: Values for the robot's locked joints, or None to keep the config's own. + + Returns: + The robot config with the overrides applied. + """ + from curobo.config_io import join_path, load_yaml + from curobo.content import get_robot_configs_path + + assert self.config.robot_config_file, "CuroboV2PlannerCfg.robot_config_file must be set." + robot_dict = load_yaml(join_path(get_robot_configs_path(), self.config.robot_config_file)) + kinematics = robot_dict["robot_cfg"]["kinematics"] + if self.config.extra_collision_spheres: + kinematics["extra_collision_spheres"] = { + **(kinematics.get("extra_collision_spheres") or {}), + **self.config.extra_collision_spheres, + } + if lock_joints: + kinematics["lock_joints"] = dict(lock_joints) + return robot_dict + + def _build_gripper_kinematics(self) -> dict[bool, Any] | None: + """Build the robot kinematics for the open and closed gripper poses. + + cuRobo folds locked-joint values into the robot's fixed transforms when a config loads, + so changing the gripper width means loading a separate kinematics config. Both are built + once here, since building one per plan would re-parse the robot description each time. + + Returns: + Kinematics keyed by whether the gripper is closed, or None if the config does not + define both poses, in which case the gripper keeps its configured width. + """ + if not (self.config.gripper_open_positions and self.config.gripper_closed_positions): + return None + + from curobo._src.types.robot import RobotCfg + + device_cfg = self.motion_planner.device_cfg + variants = {False: self.config.gripper_open_positions, True: self.config.gripper_closed_positions} + return { + is_closed: RobotCfg.create( + self._load_robot_config_dict(lock_joints=positions), device_cfg, num_envs=1 + ).kinematics.kinematics_config + for is_closed, positions in variants.items() + } + + def _set_gripper_state(self, is_closed: bool) -> None: + """Set the gripper width the planner collision-checks against. + + A gripper holding an object is closed around it, so planning with the fingers modeled + open overstates the gripper's width and rejects valid plans. The inverse kinematics, + trajectory optimization, and graph planner share one kinematics config, so this update + reaches all three. + + Call before attaching an object: the update overwrites the link spheres that hold the + fitted attachment, so calling it afterwards would discard that attachment. + + Args: + is_closed: Whether the gripper is holding an object. + """ + if self._gripper_kinematics is None or self._gripper_is_closed == is_closed: + return + self.motion_planner.kinematics.update_kinematics_config(self._gripper_kinematics[is_closed]) + self._gripper_is_closed = is_closed + + def _compute_active_joint_names(self) -> list[str]: + """Return the planner's joint names with any locked joints removed.""" + kinematics_config = self.motion_planner.kinematics.config.kinematics_config + locked_state = getattr(kinematics_config, "lock_jointstate", None) + locked_names = set(getattr(locked_state, "joint_names", None) or []) + return [name for name in self.motion_planner.joint_names if name not in locked_names] + + # ------------------------------------------------------------------ + # Frame conversion + # ------------------------------------------------------------------ + + def _refresh_base_frame_transform(self) -> None: + """Re-read the robot base pose and cache the transforms derived from it. + + cuRobo plans in the robot base frame while callers exchange poses in the env-relative + frame. Caching both directions here keeps the conversion off the per-pose path. + """ + base_pose_in_env = self.datastream.get_robot_root_pose(env_ids=[self.env_id])[0] + identity = torch.eye(4, dtype=base_pose_in_env.dtype, device=base_pose_in_env.device) + self._base_pose_in_env = base_pose_in_env + self._env_to_base_transform = torch.linalg.inv(base_pose_in_env) + self._base_frame_is_identity = bool(torch.allclose(base_pose_in_env, identity, atol=1e-6)) + + def _pose_env_to_base(self, pose_env: torch.Tensor) -> torch.Tensor: + """Express an env-relative pose in the robot base frame. + + Returns the pose unchanged when the robot base sits at the env origin. + + Args: + pose_env: Pose as a 4x4 transformation matrix in the env-relative frame. + + Returns: + The same pose as a 4x4 transformation matrix in the robot base frame. + """ + if self._base_frame_is_identity or self._env_to_base_transform is None: + return pose_env + return self._env_to_base_transform.to(device=pose_env.device, dtype=pose_env.dtype) @ pose_env + + def _pose_base_to_env(self, pose_base: torch.Tensor) -> torch.Tensor: + """Express a robot-base-frame pose in the env-relative frame. + + Inverse of :meth:`_pose_env_to_base`, applied to planned poses before returning them. + + Args: + pose_base: Pose as a 4x4 transformation matrix in the robot base frame. + + Returns: + The same pose as a 4x4 transformation matrix in the env-relative frame. + """ + if self._base_frame_is_identity or self._base_pose_in_env is None: + return pose_base + return self._base_pose_in_env.to(device=pose_base.device, dtype=pose_base.dtype) @ pose_base + + def _pose_matrix_to_curobo(self, pose_mat: torch.Tensor) -> Pose: + """Convert a 4x4 transformation matrix into a cuRobo :class:`Pose`. + + Reorders the quaternion from Isaac Lab's ``(x, y, z, w)`` to cuRobo's ``(w, x, y, z)``. + + Args: + pose_mat: Pose as a 4x4 transformation matrix in the robot base frame. + + Returns: + The equivalent batched (``[1, ...]``) cuRobo pose. + """ + import isaaclab.utils.math as PoseUtils # deferred so the module imports without a sim + + position_xyz, rot_mat = PoseUtils.unmake_pose(pose_mat) + quat_xyzw = PoseUtils.quat_from_matrix(rot_mat) + return self._make_pose( + position_xyz=position_xyz.unsqueeze(0), + quaternion_wxyz=torch.roll(quat_xyzw, shifts=1, dims=-1).unsqueeze(0), + ) + + # ------------------------------------------------------------------ + # Planning + # ------------------------------------------------------------------ + + def update_world_and_plan_motion( + self, + target_pose: torch.Tensor, + expected_attached_object: str | None = None, + env_id: int = 0, + **kwargs: Any, + ) -> bool: + """Sync the collision world, handle attachments, and plan to ``target_pose``. + + Args: + target_pose: Target end-effector pose as a 4x4 transformation matrix [m, rad] in the + env-relative frame. + expected_attached_object: Object the gripper is holding, which is collision-checked + as part of the robot. Resolved against the obstacles found on the USD stage at + construction. ``None`` plans with nothing attached. + env_id: Environment index, which must be the env this planner serves. Every other + world read uses :attr:`env_id`, so a mismatch would mix state across envs. + + Returns: + ``True`` if planning succeeded and a trajectory is now stored, else ``False``. + """ + del kwargs + assert env_id == self.env_id, ( + f"Planner for env {self.env_id} was asked to plan for env {env_id}. Every other world" + " read in this call uses self.env_id, so the two must agree." + ) + + # Trajectory optimization, attachment, and forward kinematics all build tensors that + # track gradients. The env loop runs under ``torch.inference_mode()``, which forbids + # that, so lift it for the duration of the call. + with torch.inference_mode(False), torch.enable_grad(): + self.reset_plan() + + # Move obstacles to their current poses. Updating them individually, rather than + # rebuilding the scene, preserves attachment and obstacle-disable state. + self._refresh_base_frame_transform() + self._sync_obstacle_poses() + + current_state = self._get_current_joint_state(env_id=env_id) + + # Close the gripper before attaching: the gripper update rewrites the link spheres + # that the attachment is fitted into. The release runs in the finally block — the + # spheres are fitted at this call's grasp pose, so surviving into the next call (even + # on an exception) would collision-check the object where it no longer is. + self._set_gripper_state(is_closed=expected_attached_object is not None) + self._update_attachment(expected_attached_object, current_state) + try: + goal_pose = self._pose_matrix_to_curobo(self._pose_env_to_base(target_pose)) + tool_frame = self.motion_planner.tool_frames[0] + full_trajectory = self._plan_three_phase( + current_state=current_state, + goal_pose=goal_pose, + tool_frame=tool_frame, + ) + + if full_trajectory is not None: + self._planned_joint_trajectory = full_trajectory + self._planned_eef_poses = self._joint_trajectory_to_eef_poses(full_trajectory) + if self.plan_visualizer is not None: + try: + self._visualize_plan(target_pose=target_pose, current_state=current_state) + except Exception as exc: # noqa: BLE001 (visualization must not break planning) + self._logger.warning("Plan visualization failed: %s", exc) + finally: + self._update_attachment(None, current_state) + + if full_trajectory is None: + if self.plan_visualizer is not None: + with contextlib.suppress(Exception): + self.plan_visualizer.mark_idle() + return False + return True + + # ------------------------------------------------------------------ + # Plan visualization + # ------------------------------------------------------------------ + + def _visualize_plan(self, target_pose: torch.Tensor, current_state: JointState) -> None: + """Send the stored plan to the Rerun visualizer. + + Draws the end-effector path, the goal, the obstacles, and the robot's collision spheres + split into the robot's own and those of the object it holds. Everything is drawn in the + robot base frame, the frame cuRobo collision-checks in, so the spheres, obstacles, and + goal line up. + + Args: + target_pose: Goal end-effector pose as a 4x4 matrix in the env-relative frame. + current_state: Joint state the collision spheres are evaluated at. + """ + with torch.inference_mode(False), torch.enable_grad(): + active_q = current_state.position + if active_q.ndim == 1: + active_q = active_q.unsqueeze(0) + # Identify the held object's spheres by the attached link's sphere indices rather + # than by position in the list. ``filter_valid=False`` keeps absolute indices so they + # line up with those, leaving inactive (radius <= 0) slots to drop here. + spheres = self.motion_planner.kinematics.get_robot_as_spheres(active_q.contiguous(), filter_valid=False)[0] + attached_idx: set[int] = set() + try: + _idx = self._attachment_manager().kinematics_params.get_sphere_index_from_link_name( + self.config.attached_object_link_name + ) + attached_idx = {int(j) for j in _idx.detach().cpu().tolist()} + except Exception: # noqa: BLE001 + attached_idx = set() + robot_spheres = [s for i, s in enumerate(spheres) if i not in attached_idx and float(s.radius) > 0.0] + attached_spheres = [s for i, s in enumerate(spheres) if i in attached_idx and float(s.radius) > 0.0] + + ee_positions = None + planned_poses = [self._pose_env_to_base(pose) for pose in self._planned_eef_poses] + if planned_poses: + ee_positions = np.array([p.detach().cpu().numpy().reshape(4, 4)[:3, 3] for p in planned_poses]) + plan_active = self._planned_trajectory_active() + world_scene = self._build_world_trimesh_scene() + + self.plan_visualizer.visualize_plan( + plan=plan_active, + target_pose=self._pose_env_to_base(target_pose), + robot_spheres=robot_spheres, + attached_spheres=attached_spheres, + ee_positions=ee_positions, + world_scene=world_scene, + ) + if ee_positions is not None: + self.plan_visualizer.animate_plan(ee_positions) + self.plan_visualizer.animate_spheres_along_path(plan=plan_active, robot_sphere_count=len(robot_spheres)) + + def _planned_trajectory_active(self) -> JointState: + """Return the stored plan reduced to the planner's active joints.""" + traj = self._planned_joint_trajectory + assert traj is not None, "No plan stored; a successful plan is required." + full_dof = traj.position.shape[-1] + flat = traj.position.reshape(-1, full_dof).contiguous() + names = list(traj.joint_names) if traj.joint_names else list(self.motion_planner.joint_names) + active = JointState.from_position(flat, joint_names=names).reorder(self._active_joint_names) + if not active.position.is_contiguous(): + active = JointState.from_position(active.position.contiguous(), joint_names=active.joint_names) + return active + + def _build_world_trimesh_scene(self): + """Build a :class:`trimesh.Scene` of the obstacles the planner collides against. + + Triangles come from the scene model, so concave shapes such as a sorting bin are drawn + with their true geometry rather than a bounding box. Each obstacle is placed at the pose + held in the collision data, which is kept current each plan, so moving objects follow. + Together these reproduce the geometry cuRobo collision-checks against. + + Returns: + The scene, or None if it cannot be built, so visualization never breaks planning. + """ + try: + import trimesh + + checker = getattr(self.motion_planner, "scene_collision_checker", None) + scene_model = getattr(checker, "scene_model", None) if checker is not None else None + data = getattr(checker, "data", None) if checker is not None else None + if data is None: + return None + + # Current pose and bounding-box dimensions per obstacle, from the collision data. + e = self._COLLISION_ENV_IDX + world_pose: dict[str, Any] = {} + world_dims: dict[str, Any] = {} + for attr in ("cuboids", "meshes"): + arr = getattr(data, attr, None) + if arr is None: + continue + for i in range(int(arr.count[e].item())): + if int(arr.enable[e, i].item()) == 0: + continue + name = str(arr.names[e][i]) + inv_pose = arr.inv_pose[e, i, :7].detach().cpu().numpy() # [x, y, z, qw, qx, qy, qz] + world_pose[name] = np.linalg.inv(self._pose_vec_to_matrix(inv_pose)) + world_dims[name] = arr.dims[e, i, :3].detach().cpu().numpy() + + scene = trimesh.Scene() + # Mesh obstacles, drawn from their own triangles at their current pose. + for mesh_obs in (getattr(scene_model, "mesh", None) or []) if scene_model is not None else []: + name = str(getattr(mesh_obs, "name", "")) + verts, faces = getattr(mesh_obs, "vertices", None), getattr(mesh_obs, "faces", None) + if verts is None or faces is None: + continue + verts = np.asarray(verts, dtype=float).reshape(-1, 3) + faces = np.asarray(faces, dtype=np.int64).reshape(-1, 3) + if verts.size == 0 or faces.size == 0: + continue + scene.add_geometry( + trimesh.Trimesh(vertices=verts, faces=faces, process=False), + node_name=name.replace("/", "_"), + transform=world_pose.get(name, np.eye(4)), + ) + # Cuboid obstacles, drawn as boxes from their dimensions. + for cub in (getattr(scene_model, "cuboid", None) or []) if scene_model is not None else []: + name = str(getattr(cub, "name", "")) + dims = world_dims.get(name) + if dims is None or float(np.min(dims)) <= 0.0: + continue + scene.add_geometry( + trimesh.creation.box(extents=dims.tolist()), + node_name=name.replace("/", "_"), + transform=world_pose.get(name, np.eye(4)), + ) + return scene if len(scene.geometry) else None + except Exception as exc: # noqa: BLE001 + self._logger.debug("world scene build for visualization failed: %s", exc) + return None + + @staticmethod + def _pose_vec_to_matrix(pose_vec: np.ndarray) -> np.ndarray: + """Convert a ``[x, y, z, qw, qx, qy, qz]`` pose vector to a 4x4 homogeneous matrix.""" + x, y, z, qw, qx, qy, qz = (float(v) for v in pose_vec[:7]) + norm = (qw * qw + qx * qx + qy * qy + qz * qz) ** 0.5 + if norm > 0: + qw, qx, qy, qz = qw / norm, qx / norm, qy / norm, qz / norm + mat = np.eye(4, dtype=float) + mat[0, 0] = 1 - 2 * (qy * qy + qz * qz) + mat[0, 1] = 2 * (qx * qy - qz * qw) + mat[0, 2] = 2 * (qx * qz + qy * qw) + mat[1, 0] = 2 * (qx * qy + qz * qw) + mat[1, 1] = 1 - 2 * (qx * qx + qz * qz) + mat[1, 2] = 2 * (qy * qz - qx * qw) + mat[2, 0] = 2 * (qx * qz - qy * qw) + mat[2, 1] = 2 * (qy * qz + qx * qw) + mat[2, 2] = 1 - 2 * (qx * qx + qy * qy) + mat[:3, 3] = (x, y, z) + return mat + + def _plan_three_phase( + self, + current_state: JointState, + goal_pose: Pose, + tool_frame: str, + ) -> JointState | None: + """Plan retreat, approach, and goal as three chained :meth:`plan_pose` calls. + + Each phase starts from the last waypoint of the one before. Retreat backs the gripper + away from whatever it currently touches, approach crosses free space to a pose short of + the goal, and goal closes the remaining distance. Retreat and goal expect contact, so the + gripper links are collision-disabled for them. Phases with zero distance are skipped. + + Args: + current_state: Joint state the first phase starts from. + goal_pose: Target end-effector pose in the robot base frame. + tool_frame: Robot frame the goal pose applies to. + + Returns: + The concatenated trajectory, or None if any phase failed. + """ + retreat_dist = float(self.config.retreat_distance or 0.0) + approach_dist = float(self.config.approach_distance or 0.0) + + phases: list[tuple[str, GoalToolPose, bool]] = [] + + if retreat_dist > 0: + ee_pose_cu = self._eef_pose_from_state(current_state) + retreat_pose = ee_pose_cu.multiply(self._local_offset_pose(0.0, 0.0, -retreat_dist)) + phases.append(("retreat", GoalToolPose.from_poses({tool_frame: retreat_pose}), True)) + + if approach_dist > 0: + approach_pose = goal_pose.multiply(self._local_offset_pose(0.0, 0.0, -approach_dist)) + phases.append(("approach", GoalToolPose.from_poses({tool_frame: approach_pose}), False)) + + phases.append(("goal", GoalToolPose.from_poses({tool_frame: goal_pose}), True)) + + # Links to ignore during contact phases. A held object is included so it does not + # register a collision against the surface it is being placed on. + base_disable_links = list(self.config.contact_disable_collision_links) + if self._currently_attached is not None and self.config.attached_object_link_name: + base_disable_links.append(self.config.attached_object_link_name) + + phase_trajectories: list[JointState] = [] + state = current_state + for name, goal_tool_poses, contact in phases: + phase_traj = self._plan_single_phase( + name=name, + goal_tool_poses=goal_tool_poses, + start_state=state, + contact=contact, + disable_links=base_disable_links, + ) + if phase_traj is None: + return None + phase_trajectories.append(phase_traj) + # Start the next phase where this one ended. Trajectory positions can carry extra + # leading dimensions, so flatten to ``[N, dof]`` and take the last row. The result + # holds every joint, while :meth:`plan_pose` accepts only the active ones. + full_dof = phase_traj.position.shape[-1] + last_pos_full = phase_traj.position.reshape(-1, full_dof)[-1:].contiguous() # [1, full_dof] + state_full = JointState.from_position(last_pos_full, joint_names=phase_traj.joint_names) + state = state_full.reorder(self._active_joint_names) + if not state.position.is_contiguous(): + state = JointState.from_position(state.position.contiguous(), joint_names=state.joint_names) + + return self._concat_phase_trajectories(phase_trajectories) + + def _plan_single_phase( + self, + name: str, + goal_tool_poses: GoalToolPose, + start_state: JointState, + contact: bool, + disable_links: list[str], + ) -> JointState | None: + """Run :meth:`plan_pose` for one phase, disabling contact links around the call. + + Args: + name: Phase name, used in failure messages. + goal_tool_poses: Goal pose for the phase. + start_state: Joint state the phase starts from. + contact: Whether the phase ends in contact, allowing collisions on ``disable_links``. + disable_links: Robot links whose collisions are ignored during a contact phase. + + Returns: + The trajectory for this phase, or None if planning failed. + """ + toggled = contact and bool(disable_links) + hand_links: list[str] = [] + saved_attached = None + if toggled: + attach_link = self.config.attached_object_link_name + # Re-enabling a link restores its spheres from the robot config, where the attached + # link has none. Doing that would erase the fitted object, so its spheres are saved + # and restored here instead of going through enable/disable_link_collision. + hand_links = [link for link in disable_links if link != attach_link] + if hand_links: + self.motion_planner.disable_link_collision(hand_links) + if attach_link in disable_links: + saved_attached = self._save_disable_attached_spheres(attach_link) + try: + result = self.motion_planner.plan_pose(goal_tool_poses, start_state) + finally: + if toggled: + if hand_links: + self.motion_planner.enable_link_collision(hand_links) + if saved_attached is not None: + self._restore_attached_spheres(*saved_attached) + + if result is None or getattr(result, "success", None) is None: + self._logger.warning("Phase %r failed: no result returned.", name) + return None + if not bool(result.success.any().item()): + self._logger.warning( + "Phase %r failed: every seed failed (status=%s).", name, getattr(result, "status", None) + ) + return None + + traj = result.interpolated_trajectory + if traj is None: + traj = result.js_solution + if traj is None or traj.position is None: + self._logger.warning("Phase %r failed: the returned trajectory was empty.", name) + return None + + last_tstep = self._extract_last_tstep_from(getattr(result, "interpolated_last_tstep", None)) + trimmed = self._trim_inclusive(traj, last_tstep) + if trimmed is None or trimmed.position is None or trimmed.position.shape[-2] == 0: + return None + return trimmed + + def _save_disable_attached_spheres(self, link_name: str): + """Disable the attached link's spheres by clearing their radii, keeping a copy. + + Args: + link_name: Link holding the attached object's spheres. + + Returns: + The sphere indices and their saved values, for :meth:`_restore_attached_spheres`. + """ + kp = self._attachment_manager().kinematics_params + idx = kp.get_sphere_index_from_link_name(link_name) + saved = kp.link_spheres[:, idx, :].clone() + kp.link_spheres[:, idx, 3] = -100.0 + return idx, saved + + def _restore_attached_spheres(self, idx, saved) -> None: + """Restore the spheres saved by :meth:`_save_disable_attached_spheres`.""" + self._attachment_manager().kinematics_params.link_spheres[:, idx, :] = saved + + def _eef_pose_from_state(self, state: JointState) -> Pose: + """Compute the end-effector pose for a ``[batch, dof]`` joint state.""" + position = state.position + if position.ndim == 2: # Add the horizon dimension compute_kinematics expects. + position = position.unsqueeze(1) + js = JointState.from_position(position, joint_names=state.joint_names) + kin_state = self.motion_planner.compute_kinematics(js) + link_pose = kin_state.tool_poses.get_link_pose(self.motion_planner.tool_frames[0]) + return self._make_pose( + position_xyz=link_pose.position, + quaternion_wxyz=link_pose.quaternion, + ) + + def _local_offset_pose(self, x: float, y: float, z: float) -> Pose: + """Build a :class:`Pose` translating by ``(x, y, z)`` [m] with no rotation.""" + return self._make_pose( + position_xyz=torch.tensor([[x, y, z]]), + ) + + def _make_pose( + self, + position_xyz: torch.Tensor, + quaternion_wxyz: torch.Tensor | None = None, + ) -> Pose: + """Build a cuRobo :class:`Pose` with contiguous float32 tensors on the planner's device. + + cuRobo's pose kernels reject non-contiguous tensors and tensors of the wrong dtype, both + of which arise from slicing and from autograd outputs. Normalizing here keeps that + handling in one place. The quaternion defaults to identity. + """ + device = self.motion_planner.device_cfg.device + position = position_xyz.to(device=device, dtype=torch.float32).contiguous() + if quaternion_wxyz is None: + return Pose(position=position) + quaternion = quaternion_wxyz.to(device=device, dtype=torch.float32).contiguous() + return Pose(position=position, quaternion=quaternion) + + @staticmethod + def _extract_last_tstep_from(last_tstep_tensor) -> int | None: + """Read a scalar final time step out of a per-batch tensor, or None if unavailable.""" + if last_tstep_tensor is None: + return None + try: + return int(last_tstep_tensor.flatten()[0].item()) + except (RuntimeError, ValueError, IndexError): + return None + + @staticmethod + def _trim_inclusive(traj: JointState, last_tstep: int | None) -> JointState | None: + """Drop the padding an interpolated trajectory carries past ``last_tstep``.""" + if traj is None or traj.position is None: + return None + total = traj.position.shape[-2] + if last_tstep is None or last_tstep <= 0: + return traj + end_idx = min(last_tstep + 1, total) + if end_idx >= total: + return traj + from curobo._src.state.state_joint_trajectory_ops import trim_joint_state_trajectory + + return trim_joint_state_trajectory(traj, start_idx=0, end_idx=end_idx) + + @staticmethod + def _concat_phase_trajectories(phases: list[JointState]) -> JointState: + """Join per-phase trajectories into one along the time dimension.""" + if len(phases) == 1: + return phases[0] + position = torch.cat([p.position for p in phases], dim=-2) + velocity = ( + torch.cat([p.velocity for p in phases], dim=-2) if all(p.velocity is not None for p in phases) else None + ) + acceleration = ( + torch.cat([p.acceleration for p in phases], dim=-2) + if all(p.acceleration is not None for p in phases) + else None + ) + jerk = torch.cat([p.jerk for p in phases], dim=-2) if all(p.jerk is not None for p in phases) else None + head = phases[0] + return JointState( + position=position, + velocity=velocity if velocity is not None else position * 0.0, + acceleration=acceleration if acceleration is not None else position * 0.0, + jerk=jerk if jerk is not None else position * 0.0, + joint_names=head.joint_names, + ) + + def get_planned_poses(self) -> list[torch.Tensor]: + """Return the plan as 4x4 end-effector pose matrices [m, rad] in the env-relative frame.""" + return list(self._planned_eef_poses) + + def has_next_waypoint(self) -> bool: + """Return whether the stored plan has waypoints left to execute.""" + return self._waypoint_index < len(self._planned_eef_poses) + + def get_next_waypoint_ee_pose(self) -> torch.Tensor: + """Return the next waypoint's end-effector pose as a 4x4 matrix and advance the iterator.""" + assert self.has_next_waypoint(), "No more waypoints in the current plan." + pose = self._planned_eef_poses[self._waypoint_index] + self._waypoint_index += 1 + return pose + + def reset_plan(self) -> None: + """Discard the stored plan and rewind the waypoint iterator.""" + self._planned_joint_trajectory = None + self._planned_eef_poses = [] + self._waypoint_index = 0 + + # ------------------------------------------------------------------ + # Collision world + # ------------------------------------------------------------------ + + def _initialize_static_world(self) -> None: + """Build the collision world from the USD stage. + + Walks the env's subtree, expresses each obstacle relative to the robot base, and skips + the prims named in ``world_ignore_substrings``. The world is built once here; later + plans only move obstacles within it, which keeps attachment and obstacle-disable state + intact between plans. + """ + if getattr(self.motion_planner, "scene_collision_checker", None) is None: + self._logger.warning( + "MotionPlanner has no scene collision checker; skipping world setup. " + "Set CuroboV2PlannerCfg.collision_cache or scene_model." + ) + return + + env_prim = self.datastream.get_env_prim_path(self.env_id) + robot_prim = self.datastream.get_robot_prim_path(self.env_id) + ignore = list(self.config.world_ignore_substrings) + + parser = UsdSceneParser() + parser.load_stage(self.datastream.get_usd_stage()) + scene = parser.get_obstacles_from_stage( + only_paths=[env_prim], + reference_prim_path=robot_prim, + ignore_substring=ignore, + ) + if self.config.obstacle_representation == "obb": + scene = self._scene_as_cuboids(scene) + else: + scene = scene.get_collision_check_world() + self.motion_planner.update_world(scene) + + checker = self.motion_planner.scene_collision_checker + self._world_obstacle_names = list(checker.get_obstacle_names(env_idx=self._COLLISION_ENV_IDX)) + self._object_mapping = self._discover_object_mapping(self._world_obstacle_names) + + moving = [k for k in self._object_mapping if not self._is_static_object(k)] + fixed = [k for k in self._object_mapping if self._is_static_object(k)] + self._logger.info( + "Found %d obstacles under %s (relative to %s); moving=%s fixed=%s", + len(self._world_obstacle_names), + env_prim, + robot_prim, + moving, + fixed, + ) + + def _scene_as_cuboids(self, scene) -> Scene: + """Replace every mesh in ``scene`` with its bounding box, leaving cuboids untouched. + + Box collision is exact for box-shaped obstacles and cheaper than a mesh query. Any other + shape becomes a conservative enclosure, so concave obstacles such as a bin are better + served by the mesh representation. + + Args: + scene: Scene of obstacles read from the stage. + + Returns: + A scene containing only cuboids. + """ + cuboids: list[Cuboid] = list(scene.cuboid or []) + for mesh in scene.mesh or []: + verts = np.asarray(mesh.vertices, dtype=np.float64) + if verts.size == 0: + continue + scale = mesh.scale + if scale is not None: + verts = verts * np.asarray(scale, dtype=np.float64).reshape(1, -1) + lo, hi = verts.min(axis=0), verts.max(axis=0) + dims = (hi - lo).tolist() + center_local = ((lo + hi) / 2.0).tolist() + # Offset the mesh's own pose by the box center to place the cuboid. + mesh_pose = Pose.from_list(list(mesh.pose), self.motion_planner.device_cfg) + center_offset = self._make_pose(position_xyz=torch.tensor([center_local])) + cuboid_pose = mesh_pose.multiply(center_offset).tolist() + cuboids.append(Cuboid(name=str(mesh.name), dims=dims, pose=cuboid_pose)) + return Scene(cuboid=cuboids) + + def _is_static_object(self, name: str) -> bool: + """Return whether ``name`` is configured as a fixed obstacle that is never moved.""" + return any(s in name.lower() for s in self._static_object_substrings) + + def _discover_object_mapping(self, world_obstacle_names: list[str]) -> dict[str, str]: + """Relate scene object names to the obstacle names cuRobo assigned them. + + cuRobo names obstacles after their USD prim path while the scene keys objects by short + name, so the two are matched on substring. Only objects the scene reports poses for can + be matched; geometry that exists solely on the stage stays where it was loaded. + + Args: + world_obstacle_names: Obstacle names present in the collision world. + + Returns: + Map of scene object name to cuRobo obstacle name. + """ + scene_object_names = list(self.datastream.get_object_poses(env_ids=[self.env_id]).keys()) + mapping: dict[str, str] = {} + for obj_name in scene_object_names: + key = obj_name.lower().replace("_", "") + # A short name can be a substring of a longer one (cube_1 vs cube_10), so prefer a + # path with a segment equal to the key over a bare substring hit. + candidates = [p for p in world_obstacle_names if key in str(p).lower().replace("_", "")] + if not candidates: + continue + exact = [p for p in candidates if key in str(p).lower().replace("_", "").split("/")] + if len(candidates) > 1 and not exact: + self._logger.warning( + "Object %r matches several obstacles %s; using %s.", obj_name, candidates, candidates[0] + ) + mapping[obj_name] = exact[0] if exact else candidates[0] + return mapping + + def _sync_obstacle_poses(self) -> None: + """Move every non-fixed obstacle to the pose its scene object currently has.""" + checker = getattr(self.motion_planner, "scene_collision_checker", None) + if checker is None or not self._object_mapping: + return + + object_poses = self.datastream.get_object_poses(env_ids=[self.env_id]) + for obj_name, curobo_name in self._object_mapping.items(): + if self._is_static_object(obj_name): + continue + pose_mat = object_poses.get(obj_name) + if pose_mat is None: + continue + self._set_obstacle_pose(curobo_name, self._pose_matrix_to_curobo(self._pose_env_to_base(pose_mat[0]))) + + # ------------------------------------------------------------------ + # Obstacle updates + # ------------------------------------------------------------------ + # ``SceneCollisionChecker.update_obstacle_pose`` and ``enable_obstacle`` search the cuboid + # list first and raise before reaching the mesh list, so a mesh obstacle cannot be updated + # through them. The methods below find the array holding each obstacle and update it there. + + def _obstacle_storage(self, curobo_name: str): + """Return the collision-data array holding ``curobo_name``, or None if absent.""" + cached = self._obstacle_storage_cache.get(curobo_name) + if cached is not None: + return cached + data = self.motion_planner.scene_collision_checker.data + for attr in ("cuboids", "meshes", "voxels"): + arr = getattr(data, attr, None) + if arr is not None and curobo_name in arr.get_names(self._COLLISION_ENV_IDX): + self._obstacle_storage_cache[curobo_name] = arr + return arr + return None + + def _set_obstacle_pose(self, curobo_name: str, w_obj_pose: Pose) -> None: + """Move one obstacle to ``w_obj_pose``.""" + arr = self._obstacle_storage(curobo_name) + if arr is None: + self._logger.warning("Obstacle %r is not in the collision world; skipping its pose update.", curobo_name) + return + arr.update_pose(curobo_name, w_obj_pose=w_obj_pose, env_idx=self._COLLISION_ENV_IDX) + + def _set_obstacle_enabled(self, curobo_name: str, enabled: bool) -> None: + """Include or exclude one obstacle from collision checking.""" + arr = self._obstacle_storage(curobo_name) + if arr is None: + self._logger.warning("Obstacle %r is not in the collision world; skipping enable=%s.", curobo_name, enabled) + return + arr.set_enabled(curobo_name, enabled, self._COLLISION_ENV_IDX) + + # ------------------------------------------------------------------ + # Attachment + # ------------------------------------------------------------------ + + def _attachment_manager(self): + """Return the attachment manager. + + The ``MotionPlanner.attachment_manager`` property looks the manager up at a path where + it does not live, so fall back to its actual location when the property fails. + """ + try: + mgr = self.motion_planner.attachment_manager + if mgr is not None: + return mgr + except AttributeError: + pass + return self.motion_planner.trajopt_solver.core.attachment_manager + + def _update_attachment(self, expected: str | None, current_state: JointState) -> None: + """Attach ``expected`` to the robot, releasing whatever was held before. + + Args: + expected: Object the gripper holds, or None to release without attaching. + current_state: Joint state the object's spheres are fitted at. + """ + if expected == self._currently_attached: + return + + # Release the previous object and restore it as a world obstacle. cuRobo's own re-enable + # path cannot reach mesh obstacles, so it is skipped in favour of _set_obstacle_enabled. + if self._currently_attached is not None: + self._attachment_manager().detach( + link_name=self.config.attached_object_link_name, + enable_obstacle_names=None, + ) + if self._attached_curobo_name is not None: + self._set_obstacle_enabled(self._attached_curobo_name, True) + self._currently_attached = None + self._attached_curobo_name = None + + if expected is None: + return + + curobo_name = self._object_mapping.get(expected) + if curobo_name is None: + self._logger.warning("Object %r is not in the collision world; planning unattached.", expected) + return + + checker = self.motion_planner.scene_collision_checker + obstacle = checker.scene_model.get_obstacle(curobo_name) if checker.scene_model is not None else None + if obstacle is None: + self._logger.warning( + "Object %r (%s) is not in the scene model; planning unattached.", expected, curobo_name + ) + return + + # Fitting spheres to an obstacle bakes in the pose the obstacle carries, and the scene + # model keeps the pose the object was loaded at rather than where it is now. Passing the + # offset between the two cancels the stale pose so the spheres land on the object. + world_pose_offset = None + current_world_pose = self._object_world_pose(expected) + if current_world_pose is not None: + obstacle_pose = Pose.from_list(list(obstacle.pose), self.motion_planner.device_cfg) + world_pose_offset = current_world_pose.multiply(obstacle_pose.inverse()) + + try: + self._attachment_manager().attach( + joint_states=current_state, + obstacles=[obstacle], + link_name=self.config.attached_object_link_name, + num_spheres=self.config.attached_object_num_spheres, + surface_radius=self.config.surface_sphere_radius, + sphere_fit_type=self._resolve_sphere_fit_type(), + world_objects_pose_offset=world_pose_offset, + disable_obstacle_names=None, + ) + except Exception as exc: # noqa: BLE001 (a failed attach should not abort planning) + self._logger.warning("Attaching %r failed: %s. Planning unattached.", expected, exc) + return + + # Record the attachment before any follow-up work: once attach() has written spheres to + # the attached link, the release branch above must run on the next call even if the + # obstacle-disable step below fails, or the spheres would be left welded to the hand. + self._currently_attached = expected + self._attached_curobo_name = curobo_name + # Exclude the object from the world so it is not counted twice, once as attached spheres + # and once as an obstacle. cuRobo's own auto-disable cannot reach mesh obstacles. + self._set_obstacle_enabled(curobo_name, False) + + def _object_world_pose(self, obj_name: str) -> Pose | None: + """Return a scene object's current pose in the robot base frame, or None if unknown.""" + pose_mat = self.datastream.get_object_poses(env_ids=[self.env_id]).get(obj_name) + if pose_mat is None: + return None + return self._pose_matrix_to_curobo(self._pose_env_to_base(pose_mat[0])) + + def _resolve_sphere_fit_type(self) -> SphereFitType: + """Convert the configured sphere-fit name into its :class:`SphereFitType` value.""" + name = (self.config.sphere_fit_type or "").upper() + try: + return SphereFitType[name] + except KeyError: + self._logger.warning( + "Unknown sphere_fit_type %r; using SURFACE. Valid names: %s", + self.config.sphere_fit_type, + [m.name for m in SphereFitType], + ) + return SphereFitType.SURFACE + + # ------------------------------------------------------------------ + # Joint state + # ------------------------------------------------------------------ + + def _get_current_joint_state(self, env_id: int) -> JointState: + """Read the robot's joint positions, keeping only the planner's active joints. + + The solvers operate on the active joints alone, so passing the full set, locked joints + included, fails downstream on a dimension mismatch. + """ + joint_pos_isaac = self.datastream.get_robot_joint_positions(env_ids=[env_id])[0] + position = ( + joint_pos_isaac.unsqueeze(0) + .to(device=self.motion_planner.device_cfg.device, dtype=torch.float32) + .contiguous() + ) + state = JointState.from_position(position, joint_names=self.datastream.get_robot_joint_names()) + reordered = state.reorder(self._active_joint_names) + # Reordering can return a view, which the CUDA kernels reject. + if not reordered.position.is_contiguous(): + reordered = JointState.from_position(reordered.position.contiguous(), joint_names=reordered.joint_names) + return reordered + + # ------------------------------------------------------------------ + # End-effector poses + # ------------------------------------------------------------------ + + def _joint_trajectory_to_eef_poses(self, traj: JointState) -> list[torch.Tensor]: + """Convert a joint trajectory into end-effector poses with one batched call. + + Args: + traj: Planned joint trajectory. + + Returns: + One 4x4 pose matrix per waypoint, in the env-relative frame. + """ + import isaaclab.utils.math as PoseUtils # deferred so the module imports without a sim + + positions = traj.position + full_dof = positions.shape[-1] + flat = positions.reshape(-1, full_dof) + + # The trajectory holds every joint, while forward kinematics takes the active ones. + traj_joint_names = traj.joint_names if traj.joint_names else self.motion_planner.joint_names + js_position = flat.unsqueeze(0).to(self.motion_planner.device_cfg.device) + full_state = JointState.from_position(js_position, joint_names=list(traj_joint_names)) + single_state = full_state.reorder(self._active_joint_names) + if not single_state.position.is_contiguous(): + single_state = JointState.from_position( + single_state.position.contiguous(), joint_names=single_state.joint_names + ) + + # Forward kinematics builds tensors that track gradients, which the env loop forbids. + with torch.inference_mode(False), torch.enable_grad(): + kin_state = self.motion_planner.compute_kinematics(single_state) + + tool_frame = self.motion_planner.tool_frames[0] + link_pose = kin_state.tool_poses.get_link_pose(tool_frame) + # Move results onto the env's device: cuRobo computes on its own CUDA device, while + # callers combine these poses with env tensors. + env_device = self.datastream.device + positions_w = link_pose.position.detach().to(env_device) # [T, 3] + quaternions_wxyz = link_pose.quaternion.detach().to(env_device) # [T, 4] + quaternions_xyzw = torch.roll(quaternions_wxyz, shifts=-1, dims=-1) + rotations = PoseUtils.matrix_from_quat(quaternions_xyzw) # [T, 3, 3] + return [ + self._pose_base_to_env(PoseUtils.make_pose(positions_w[t], rotations[t])) + for t in range(positions_w.shape[0]) + ] diff --git a/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner_cfg.py b/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner_cfg.py new file mode 100644 index 0000000..9eb46b1 --- /dev/null +++ b/isaac_autodata_interfaces/motion_planners/curobo_v2/curobo_v2_planner_cfg.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration for the cuRobo v2 motion planner. + +Named presets are built by the factory methods at the end of the class and selected either by +profile name or by task id, so an environment profile resolves the same name against any planner +backend. + +This module imports neither cuRobo nor Isaac Lab, so it can be loaded without a simulator. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class CuroboV2PlannerCfg: + """Configuration for the cuRobo v2 motion planner. + + Args: + robot_config_file: cuRobo robot config (e.g. ``"franka.yml"``), resolved through cuRobo's + robot config search path. Absolute paths are also accepted. + robot_name: Robot identifier used in logs and visualization. + gripper_open_positions: Joint positions [m] the gripper holds while moving empty-handed. + gripper_closed_positions: Joint positions [m] the gripper holds while carrying an object. + Leave this or the open positions empty to keep the robot config's own width. + hand_link_names: Gripper links, used to populate ``contact_disable_collision_links``. + attached_object_link_name: Link a grasped object is attached to. + extra_collision_spheres: Collision spheres to allocate per link, applied on top of the + robot config's own values. + scene_model: cuRobo scene config holding fixed collision geometry, or None for none. + Obstacles are read from the USD stage regardless. + static_objects: Substrings of scene object names to treat as fixed. Matching obstacles + are placed once and never moved again; everything else follows its object. + world_ignore_substrings: Substrings of USD prim paths to skip when reading obstacles. + obstacle_representation: ``"mesh"`` keeps each obstacle's triangles, so concave shapes + such as a bin are represented faithfully and the gripper can reach inside. + ``"obb"`` reduces every obstacle to a bounding box, which is exact for boxes and + cheaper, but unsuitable for concave geometry. + num_ik_seeds: Number of inverse-kinematics seeds. + num_trajopt_seeds: Number of trajectory-optimization seeds. + position_tolerance: Goal position tolerance [m]. + orientation_tolerance: Goal orientation tolerance [rad]. + optimizer_collision_activation_distance: Distance at which collision costs engage [m]. + use_cuda_graph: Whether to use CUDA graphs for speed. + random_seed: Seed for reproducible planning. + collision_cache: Obstacle slots to pre-allocate per type. Required when no + ``scene_model`` is given so obstacles can be loaded later. + motion_noise_scale: Action noise applied per waypoint during execution. Read by the + SkillGen algorithm, not by cuRobo. + surface_sphere_radius: Amount each fitted sphere is inflated by [m]. + sphere_fit_type: How spheres are fitted to a grasped object. One of ``"SURFACE"``, + ``"VOXEL"``, or ``"MORPHIT"``. + attached_object_num_spheres: Spheres to fit to a grasped object. Must not exceed the + attached link's allocation in ``extra_collision_spheres``, or attaching fails. + approach_distance: Distance short of the goal the approach phase stops at [m]. + retreat_distance: Distance the gripper backs off before crossing free space [m]. + contact_disable_collision_links: Links whose collisions are ignored during the phases + that end in contact, so the gripper may touch what it is reaching for. + visualize_spheres: Not implemented; use ``visualize_plan``. + visualize_plan: Draw the plan in Rerun: end-effector path, goal, obstacles, and + collision spheres. Applies to env 0 only and requires the ``rerun-sdk`` package. + """ + + # Robot + robot_config_file: str | None = None + robot_name: str = "" + + # Gripper. The open and closed positions are loaded into the robot's locked joints as the + # grasp state changes, so the gripper is collision-checked at the width it actually has. + gripper_open_positions: dict[str, float] = field(default_factory=dict) + gripper_closed_positions: dict[str, float] = field(default_factory=dict) + hand_link_names: list[str] = field(default_factory=list) + attached_object_link_name: str = "attached_object" + # The stock Franka config allocates 4 spheres to the attached link, too coarse a model of a + # grasped object to plan through the narrow clearances of a bin. + extra_collision_spheres: dict[str, int] = field(default_factory=lambda: {"attached_object": 100}) + + # Scene + scene_model: str | None = None + static_objects: list[str] = field(default_factory=list) + world_ignore_substrings: list[str] = field( + default_factory=lambda: ["/World/defaultGroundPlane", "/curobo", "/Robot"] + ) + obstacle_representation: str = "mesh" + + # Planning + num_ik_seeds: int = 32 + num_trajopt_seeds: int = 4 + position_tolerance: float = 0.005 + orientation_tolerance: float = 0.05 + optimizer_collision_activation_distance: float = 0.03 + use_cuda_graph: bool = True + random_seed: int = 123 + collision_cache: dict[str, int] = field(default_factory=lambda: {"obb": 64, "mesh": 64}) + motion_noise_scale: float = 0.0 + + # Attachment + surface_sphere_radius: float = 0.005 + sphere_fit_type: str = "SURFACE" + attached_object_num_spheres: int = 100 + + # Approach and contact + approach_distance: float = 0.05 + retreat_distance: float = 0.05 + contact_disable_collision_links: list[str] = field(default_factory=list) + + # Visualization + visualize_spheres: bool = False + visualize_plan: bool = False + + def __post_init__(self) -> None: + """Validate cross-field constraints that would otherwise fail deep inside cuRobo.""" + assert self.obstacle_representation in ( + "mesh", + "obb", + ), f"obstacle_representation must be 'mesh' or 'obb', got {self.obstacle_representation!r}" + # When the allocation comes from the robot config instead, its size is unknown here and + # is checked by cuRobo at attach time. + if self.attached_object_link_name in self.extra_collision_spheres: + allocated = self.extra_collision_spheres[self.attached_object_link_name] + assert self.attached_object_num_spheres <= allocated, ( + f"attached_object_num_spheres ({self.attached_object_num_spheres}) exceeds the " + f"{self.attached_object_link_name!r} allocation in extra_collision_spheres ({allocated})" + ) + + # ------------------------------------------------------------------ + # Factory methods + # ------------------------------------------------------------------ + + @classmethod + def franka_config(cls) -> CuroboV2PlannerCfg: + """Create a configuration for the Franka Panda arm and its parallel gripper.""" + return cls( + robot_config_file="franka.yml", + robot_name="franka", + gripper_open_positions={"panda_finger_joint1": 0.04, "panda_finger_joint2": 0.04}, + gripper_closed_positions={"panda_finger_joint1": 0.023, "panda_finger_joint2": 0.023}, + hand_link_names=["panda_leftfinger", "panda_rightfinger", "panda_hand"], + motion_noise_scale=0.02, + ) + + @classmethod + def franka_stack_cube_config(cls) -> CuroboV2PlannerCfg: + """Create a configuration for the Franka cube-stacking task. + + The table is fixed while the cubes follow their scene objects, and the gripper links are + allowed to touch during the contact phases so the fingers can close on a cube. + """ + cfg = cls.franka_config() + cfg.static_objects = ["table"] + cfg.optimizer_collision_activation_distance = 0.01 + cfg.approach_distance = 0.05 + cfg.retreat_distance = 0.05 + cfg.surface_sphere_radius = 0.01 + cfg.collision_cache = {"obb": 150, "mesh": 150} + cfg.contact_disable_collision_links = list(cfg.hand_link_names) + return cfg + + @classmethod + def franka_stack_cube_bin_config(cls) -> CuroboV2PlannerCfg: + """Create a configuration for stacking cubes inside a sorting bin. + + The bin joins the table as fixed geometry. Clearances are tighter than on an open table, + so the collision margin is wider, the retreat longer, and the modeled closed-finger width + slightly larger for margin against the walls. Obstacles keep their own triangles, letting + the planner avoid the bin walls while the gripper reaches inside. + """ + cfg = cls.franka_stack_cube_config() + cfg.static_objects = ["blue_sorting_bin", "bin", "table"] + cfg.optimizer_collision_activation_distance = 0.02 + cfg.approach_distance = 0.05 + cfg.retreat_distance = 0.07 + cfg.surface_sphere_radius = 0.01 + cfg.gripper_closed_positions = {"panda_finger_joint1": 0.024, "panda_finger_joint2": 0.024} + return cfg + + @classmethod + def from_profile(cls, profile_name: str) -> CuroboV2PlannerCfg: + """Create a configuration from a named planner profile. + + Profiles let an environment profile name the tuning that suits its scene, rather than + inferring it from the task id. + + Args: + profile_name: Key into :data:`PLANNER_PROFILES`. + + Returns: + Configuration for the named profile. + """ + assert ( + profile_name in PLANNER_PROFILES + ), f"Unknown planner profile {profile_name!r}. Registered: {sorted(PLANNER_PROFILES)}" + return PLANNER_PROFILES[profile_name]() + + @classmethod + def from_task_name(cls, task_name: str) -> CuroboV2PlannerCfg: + """Create a configuration by matching substrings of the task id. + + Args: + task_name: Task id, e.g. ``"Isaac-Stack-Cube-Franka-IK-Rel-Skillgen-v0"``. + + Returns: + Configuration for the task, falling back to the plain Franka preset. + """ + lower = task_name.lower() + if "stack-cube-bin" in lower: + return cls.franka_stack_cube_bin_config() + if "stack-cube" in lower: + return cls.franka_stack_cube_config() + logging.getLogger(__name__).warning( + "No planner preset matches task %r; falling back to the Franka configuration.", task_name + ) + return cls.franka_config() + + def to_v2_kwargs(self) -> dict[str, Any]: + """Return the fields :meth:`MotionPlannerCfg.create` accepts, ready to splat into it. + + Fields read elsewhere, such as ``motion_noise_scale`` and the visualization flags, are + left out. + """ + return { + "robot": self.robot_config_file, + "scene_model": self.scene_model, + "collision_cache": self.collision_cache, + "num_ik_seeds": self.num_ik_seeds, + "num_trajopt_seeds": self.num_trajopt_seeds, + "position_tolerance": self.position_tolerance, + "orientation_tolerance": self.orientation_tolerance, + "optimizer_collision_activation_distance": self.optimizer_collision_activation_distance, + "use_cuda_graph": self.use_cuda_graph, + "random_seed": self.random_seed, + } + + +PLANNER_PROFILES: dict[str, Callable[[], CuroboV2PlannerCfg]] = { + "franka": CuroboV2PlannerCfg.franka_config, + "franka_stack_cube": CuroboV2PlannerCfg.franka_stack_cube_config, + "franka_stack_cube_bin": CuroboV2PlannerCfg.franka_stack_cube_bin_config, +} +"""Maps a planner-profile name to the factory that builds its configuration. + +Keys match those of the other planner backends, so an environment profile's ``planner`` name +resolves whichever backend is selected. Register new profiles here.""" diff --git a/isaac_autodata_interfaces/motion_planners/curobo_v2/plan_visualizer.py b/isaac_autodata_interfaces/motion_planners/curobo_v2/plan_visualizer.py new file mode 100644 index 0000000..b5ab9fa --- /dev/null +++ b/isaac_autodata_interfaces/motion_planners/curobo_v2/plan_visualizer.py @@ -0,0 +1,581 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rerun visualizer for the cuRobo v2 motion planner. + +Draws the planned end-effector trajectory, the goal, the world obstacles, and the collision +spheres of both the robot and any object it holds. The planner supplies the two sphere lists +already separated, all in the robot base frame, so nothing here needs to transform them. + +Rerun is an optional dependency. The planner imports this module only when +:attr:`CuroboV2PlannerCfg.visualize_plan` is set, so other runs never require it. +""" + +from __future__ import annotations + +import atexit +import contextlib +import numpy as np +import os +import signal +import threading +import time +import torch +import weakref +from typing import TYPE_CHECKING, Any + + +def _import_rerun_sdk(): + """Import the Rerun SDK, working around another package that claims the same import name. + + An unrelated ``rerun`` file-watcher package installs directly into ``site-packages/rerun`` + while ``rerun-sdk`` ships under ``rerun_sdk/``, so the file-watcher can win the import and + yield a module with no ``init`` or ``log``. Putting the SDK's directory first on the search + path avoids that. + """ + import importlib + import os + import sys + + # This works because the planner imports this module lazily and nothing before it imports + # ``rerun``, so the search path still decides. An already-imported wrong module is not + # reloaded: the SDK is a native extension and re-importing it over the other package crashes + # the interpreter. + if "rerun" not in sys.modules: + for entry in list(sys.path): + candidate = os.path.join(entry, "rerun_sdk") + if os.path.isfile(os.path.join(candidate, "rerun", "__init__.py")): + sys.path.insert(0, candidate) + break + + module = importlib.import_module("rerun") + if hasattr(module, "init"): + return module + + raise ImportError( + "cuRobo v2 plan visualization needs the Rerun robotics SDK, but `import rerun` resolved to " + f"the deprecated 'rerun' file-watcher package ({getattr(module, '__file__', '?')}) — it was " + "imported before this module could prefer rerun-sdk. Fix the env with `pip uninstall rerun` " + "(this keeps rerun-sdk), or disable visualize_plan." + ) + + +rr = _import_rerun_sdk() + +_RR_HAS_TRANSFORM_AXES = hasattr(rr, "TransformAxes3D") + +try: + import psutil + + PSUTIL_AVAILABLE = True +except ImportError: + PSUTIL_AVAILABLE = False + print("Warning: psutil not available. Rerun process monitoring will be limited.") + +if TYPE_CHECKING: + import trimesh + + +# Global registry to track all visualizer instances for cleanup. +_GLOBAL_PLAN_VISUALIZERS: list[PlanVisualizer] = [] + + +def _kill_owned_rerun_processes() -> int: + """Kill Rerun viewer processes spawned by this process, and only those. + + Matching on the process tree rather than on the name alone keeps viewers belonging to other + users or jobs on a shared host untouched. + + Returns: + Number of processes killed. + """ + if not PSUTIL_AVAILABLE: + return 0 + killed = 0 + with contextlib.suppress(psutil.Error): + for proc in psutil.Process(os.getpid()).children(recursive=True): + try: + name = proc.name().lower() + cmdline = " ".join(proc.cmdline()).lower() + if "rerun" in name or "rerun" in cmdline: + proc.kill() + killed += 1 + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + pass + return killed + + +def _cleanup_all_plan_visualizers() -> None: + """Close tracked visualizers and kill this process's own Rerun viewers on exit.""" + for visualizer in _GLOBAL_PLAN_VISUALIZERS[:]: + if not visualizer._closed: + visualizer.close() + _GLOBAL_PLAN_VISUALIZERS.clear() + _kill_owned_rerun_processes() + + +atexit.register(_cleanup_all_plan_visualizers) + + +class PlanVisualizer: + """Visualizes cuRobo v2 motion plans, collision spheres, and obstacles via Rerun. + + Args: + robot_name: Robot identifier used in the recording id. + recording_id: Optional Rerun recording id; defaults to ``motion_plan_``. + debug: Whether to print debug information. + save_path: Optional path to save the Rerun recording on close. + base_translation: Optional translation added to every visualized entity. Defaults to zero + because the v2 planner works entirely in the robot-base frame. + """ + + def __init__( + self, + robot_name: str = "franka", + recording_id: str | None = None, + debug: bool = False, + save_path: str | None = None, + base_translation: np.ndarray | None = None, + ) -> None: + self.robot_name = robot_name + self.debug = debug + self.recording_id = recording_id or f"motion_plan_{robot_name}" + self.save_path = save_path + self._closed = False + self._base_translation = ( + np.array(base_translation, dtype=float) if base_translation is not None else np.zeros(3) + ) + + self._parent_pid = os.getpid() + self._monitor_thread: threading.Thread | None = None + self._monitor_active = False + + # cuRobo v2 MotionPlanner reference, set by the planner for sphere animation. + self._motion_planner_ref: Any = None + + _GLOBAL_PLAN_VISUALIZERS.append(self) + + rr.init(self.recording_id, spawn=False) + self._rerun_process = None + self._sink = self._connect_sink() + + rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP) + + self._current_frame = 0 + self._sphere_entities: dict[str, list[str]] = {"robot": [], "attached": [], "target": []} + + self._start_parent_process_monitoring() + + self._finalizer = weakref.finalize( + self, self._cleanup_class_resources, self.recording_id, self.save_path, debug, self._sink + ) + self._atexit_callback = atexit.register( + self._cleanup_class_resources, self.recording_id, self.save_path, debug, self._sink + ) + + self._original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_DFL) + self._original_sigterm_handler = signal.signal(signal.SIGTERM, signal.SIG_DFL) + + def signal_handler(signum, frame): + if self.debug: + print(f"Received signal {signum}, closing Rerun viewer...") + self._cleanup_on_exit() + if signum == signal.SIGINT: + signal.signal(signal.SIGINT, self._original_sigint_handler) + elif signum == signal.SIGTERM: + signal.signal(signal.SIGTERM, self._original_sigterm_handler) + os.kill(os.getpid(), signum) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + if self.debug: + print(f"Initialized cuRobo v2 Rerun visualization (recording id: {self.recording_id})") + + # ------------------------------------------------------------------ + # Process lifecycle + # ------------------------------------------------------------------ + + def _connect_sink(self) -> str: + """Attach a live viewer if one can be launched, otherwise record to an ``.rrd`` file. + + Spawning a viewer needs the ``rerun`` executable, which may be missing from the path, so + the SDK's bundled binary is used when present. Which route was taken is always logged, + since a viewer that fails to start is otherwise invisible. + + Returns: + ``"viewer"``, ``"file"``, or ``"none"``, describing where output went. + """ + exe = self._bundled_viewer_path() + try: + if exe is not None: + rr.spawn(executable_path=exe) + else: + rr.spawn() + print( + f"[PlanVisualizer] live Rerun viewer launched (recording_id={self.recording_id!r}, " + f"viewer={exe or 'rerun on PATH'}).", + flush=True, + ) + return "viewer" + except Exception as exc: # noqa: BLE001 + if self.save_path: + try: + rr.save(self.save_path) + print( + f"[PlanVisualizer] no live viewer ({exc}); recording to {self.save_path!r}. " + f"Open it with: rerun {self.save_path}", + flush=True, + ) + return "file" + except Exception as save_exc: # noqa: BLE001 + print(f"[PlanVisualizer] viewer spawn AND file save failed: {save_exc!r}", flush=True) + else: + print(f"[PlanVisualizer] could not launch a live viewer ({exc}); no save_path set.", flush=True) + return "none" + + @staticmethod + def _bundled_viewer_path() -> str | None: + """Path to rerun-sdk's bundled viewer binary (``rerun_cli/rerun``), or ``None`` if absent.""" + import os + + try: + pkg_dir = os.path.dirname(rr.__file__) # .../rerun_sdk/rerun + candidate = os.path.abspath(os.path.join(pkg_dir, os.pardir, "rerun_cli", "rerun")) + return candidate if os.path.isfile(candidate) else None + except Exception: # noqa: BLE001 + return None + + def _start_parent_process_monitoring(self) -> None: + if not PSUTIL_AVAILABLE: + return + self._monitor_active = True + + def monitor_parent_process() -> None: + parent_process = psutil.Process(self._parent_pid) + while self._monitor_active: + try: + if not parent_process.is_running(): + self._kill_rerun_processes() + break + time.sleep(2) + except (psutil.NoSuchProcess, psutil.AccessDenied): + self._kill_rerun_processes() + break + except Exception: + break + + self._monitor_thread = threading.Thread(target=monitor_parent_process, daemon=True) + self._monitor_thread.start() + + def _kill_rerun_processes(self) -> None: + try: + _kill_owned_rerun_processes() + except Exception as exc: # pragma: no cover + if self.debug: + print(f"Error killing rerun processes: {exc}") + + @staticmethod + def _cleanup_class_resources(recording_id: str, save_path: str | None, debug: bool, sink: str = "") -> None: + # Save before disconnecting: data logged so far is not written to a file opened after + # the recording is closed. Skip when the sink already is that file, since re-opening it + # would truncate what was streamed to it. + if save_path is not None and sink != "file": + rr.save(save_path) + rr.disconnect() + _kill_owned_rerun_processes() + + def _cleanup_on_exit(self) -> None: + if not self._closed: + self._monitor_active = False + self.close() + self._kill_rerun_processes() + + def close(self) -> None: + """Close the Rerun connection, terminate the viewer, and deregister.""" + if self._closed: + return + self._monitor_active = False + if self._monitor_thread and self._monitor_thread.is_alive(): + time.sleep(0.1) + # Save before disconnecting so the recording keeps the data logged so far; skip when the + # sink already is the file, since re-opening it would truncate it. + if self.save_path is not None and self._sink != "file": + rr.save(self.save_path) + rr.disconnect() + # The instance is now fully released; the exit-time callback would only repeat this work + # on a closed recording. + atexit.unregister(self._cleanup_class_resources) + self._finalizer.detach() + self._closed = True + try: + process = getattr(self, "_rerun_process", None) + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except Exception: + process.kill() + except Exception: + pass + self._kill_rerun_processes() + if self in _GLOBAL_PLAN_VISUALIZERS: + _GLOBAL_PLAN_VISUALIZERS.remove(self) + + def set_motion_planner_reference(self, motion_planner: Any) -> None: + """Provide the cuRobo v2 ``MotionPlanner`` used to compute spheres during animation.""" + self._motion_planner_ref = motion_planner + + # ------------------------------------------------------------------ + # Static plan visualization + # ------------------------------------------------------------------ + + def visualize_plan( + self, + plan: Any, + target_pose: torch.Tensor, + robot_spheres: list[Any] | None = None, + attached_spheres: list[Any] | None = None, + ee_positions: np.ndarray | None = None, + world_scene: trimesh.Scene | None = None, + ) -> None: + """Log one static snapshot of a plan: obstacles, target, EE trajectory, and spheres. + + Args: + plan: Active-DoF joint trajectory (only ``plan.position`` is read, for the EE path + fallback). The planner normally supplies ``ee_positions`` directly. + target_pose: Target end-effector pose as a 4x4 matrix in the robot-base frame. + robot_spheres: Robot collision spheres (cuRobo ``Sphere`` objects). + attached_spheres: Attached-object collision spheres. + ee_positions: ``[T, 3]`` end-effector positions in the robot-base frame. + world_scene: Optional ``trimesh.Scene`` of the world obstacles. + """ + rr.set_time("static_plan", sequence=self._current_frame) + self._current_frame += 1 + + self._clear_visualization() + + if world_scene is not None: + self._visualize_world_scene(world_scene) + self._visualize_target_pose(target_pose) + self._visualize_trajectory(plan, ee_positions) + + if robot_spheres: + self._log_spheres(robot_spheres, "robot", [0, 255, 100, 128]) + if attached_spheres: + self._log_spheres(attached_spheres, "attached", [255, 0, 0, 128]) + else: + self._clear_attached_spheres() + + n_ee = 0 if ee_positions is None else len(ee_positions) + print( + f"[PlanVisualizer] logged plan -> sink={self._sink}: {n_ee} EE waypoints, " + f"{len(robot_spheres or [])} robot spheres, {len(attached_spheres or [])} attached spheres.", + flush=True, + ) + + def _clear_visualization(self) -> None: + for path in ("trajectory", "target", "anim"): + rr.log(f"world/{path}", rr.Clear(recursive=True)) + for entity_type, entities in self._sphere_entities.items(): + for entity in entities: + rr.log(f"world/{entity_type}/{entity}", rr.Clear(recursive=True)) + self._sphere_entities[entity_type] = [] + + def clear_visualization(self) -> None: + """Public wrapper around :meth:`_clear_visualization`.""" + self._clear_visualization() + + def _visualize_target_pose(self, target_pose: torch.Tensor) -> None: + mat = target_pose.detach().cpu().numpy() if torch.is_tensor(target_pose) else np.asarray(target_pose) + mat = mat.reshape(4, 4) + pos = mat[:3, 3] + self._base_translation + rot = mat[:3, :3] + rr.log("world/target/position", rr.Points3D(positions=np.array([pos]), colors=[[255, 0, 0]], radii=[0.02])) + rr.log("world/target/frame", rr.Transform3D(translation=pos, mat3x3=rot)) + + def _visualize_trajectory(self, plan: Any, ee_positions: np.ndarray | None) -> None: + if ee_positions is None: + raw = plan.position.detach().cpu().numpy() if torch.is_tensor(plan.position) else np.array(plan.position) + if raw.ndim != 2 or raw.shape[1] < 3: + return # nothing sensible to draw without explicit EE positions + positions = raw[:, :3] + else: + positions = np.asarray(ee_positions) + if positions.size == 0: + return + positions = positions + self._base_translation + rr.log("world/trajectory", rr.LineStrips3D([positions], colors=[[0, 100, 255]], radii=[0.005]), static=True) + for i, pos in enumerate(positions): + rr.log( + f"world/trajectory/keyframe_{i}", + rr.Points3D(positions=np.array([pos]), colors=[[0, 100, 255]], radii=[0.01]), + static=True, + ) + + def _log_spheres(self, spheres: list[Any], entity_type: str, color: list[int]) -> None: + for i, sphere in enumerate(spheres): + entity_id = f"sphere_{i}" + self._sphere_entities.setdefault(entity_type, []).append(entity_id) + pos = ( + sphere.position.detach().cpu().numpy() + if torch.is_tensor(sphere.position) + else np.array(sphere.position) + ).reshape(-1) + pos = pos + self._base_translation + rr.log( + f"world/{entity_type}/{entity_id}", + rr.Points3D(positions=np.array([pos]), colors=[color], radii=[float(sphere.radius)]), + ) + + def _clear_attached_spheres(self) -> None: + for entity_id in self._sphere_entities.get("attached", []): + rr.log(f"world/attached/{entity_id}", rr.Clear(recursive=True)) + self._sphere_entities["attached"] = [] + + def _visualize_world_scene(self, scene: trimesh.Scene) -> None: + import trimesh + + if not hasattr(self, "_logged_geometry"): + self._logged_geometry: set[str] = set() + + for node in scene.graph.nodes_geometry: + tform, geom_key = scene.graph.get(node) + mesh = scene.geometry.get(geom_key) + if mesh is None: + continue + rr_path = f"world/scene/{node.replace('/', '_')}" + if _RR_HAS_TRANSFORM_AXES: + rr.log(rr_path, rr.Transform3D(translation=tform[:3, 3], mat3x3=tform[:3, :3]), static=False) + else: + rr.log( + rr_path, + rr.Transform3D(translation=tform[:3, 3], mat3x3=tform[:3, :3], axis_length=0.0), + static=False, + ) + if rr_path not in self._logged_geometry: + if isinstance(mesh, trimesh.Trimesh): + rr.log( + rr_path, + rr.Mesh3D( + vertex_positions=mesh.vertices, + triangle_indices=mesh.faces, + vertex_normals=mesh.vertex_normals if mesh.vertex_normals is not None else None, + ), + static=True, + ) + self._logged_geometry.add(rr_path) + + # ------------------------------------------------------------------ + # Animation + # ------------------------------------------------------------------ + + def animate_plan(self, ee_positions: np.ndarray, timeline: str = "plan", point_radius: float = 0.01) -> None: + """Play back the end-effector marker along ``ee_positions`` on ``timeline``.""" + if ee_positions is None or len(ee_positions) == 0: + return + for idx, pos in enumerate(ee_positions): + rr.set_time(timeline, sequence=idx) + rr.log( + "world/anim/ee", + rr.Points3D( + positions=np.array([pos + self._base_translation]), colors=[[0, 100, 255]], radii=[point_radius] + ), + ) + + def animate_spheres_along_path( + self, + plan: Any, + robot_sphere_count: int, + timeline: str = "sphere_animation", + interpolation_steps: int = 10, + ) -> None: + """Animate robot (green) and attached (orange) spheres along the planned trajectory. + + Recomputes collision spheres at densely interpolated configurations via the v2 + ``MotionPlanner.kinematics.get_robot_as_spheres``. ``plan`` must carry active-DoF + positions (``[T, active_dof]``); ``robot_sphere_count`` is the number of robot self + spheres (the remainder of each frame's active spheres are the attached object's). + """ + motion_planner = self._motion_planner_ref + if motion_planner is None or plan is None or len(plan.position) == 0: + return + device = motion_planner.device_cfg.device + + self._hide_static_spheres_for_animation() + interpolated = self._create_interpolated_trajectory(plan, interpolation_steps) + + for frame_idx, joint_positions in enumerate(interpolated): + rr.set_time(timeline, sequence=frame_idx) + q = joint_positions if isinstance(joint_positions, torch.Tensor) else torch.tensor(joint_positions) + q = q.to(device=device, dtype=torch.float32) + if q.ndim == 1: + q = q.unsqueeze(0) # [active_dof] -> [1, active_dof] + try: + with torch.inference_mode(False), torch.enable_grad(): + sphere_list = motion_planner.kinematics.get_robot_as_spheres(q)[0] + except Exception as exc: + if self.debug: + print(f"Failed to compute spheres for frame {frame_idx}: {exc}") + continue + + robot_pos, robot_rad, att_pos, att_rad = [], [], [], [] + for i, sphere in enumerate(sphere_list): + pos = ( + sphere.position.detach().cpu().numpy() + if torch.is_tensor(sphere.position) + else np.array(sphere.position) + ).reshape(-1) + self._base_translation + if i < robot_sphere_count: + robot_pos.append(pos) + robot_rad.append(float(sphere.radius)) + else: + att_pos.append(pos) + att_rad.append(float(sphere.radius)) + + if robot_pos: + rr.log( + "world/robot_animation", + rr.Points3D( + positions=np.array(robot_pos), colors=[[0, 255, 100, 220]] * len(robot_pos), radii=robot_rad + ), + ) + if att_pos: + rr.log( + "world/attached_animation", + rr.Points3D(positions=np.array(att_pos), colors=[[255, 150, 0, 220]] * len(att_pos), radii=att_rad), + ) + else: + rr.log("world/attached_animation", rr.Clear(recursive=True)) + + def _hide_static_spheres_for_animation(self) -> None: + for entity_id in self._sphere_entities.get("robot", []): + rr.log(f"world/robot/{entity_id}", rr.Clear(recursive=True)) + for entity_id in self._sphere_entities.get("attached", []): + rr.log(f"world/attached/{entity_id}", rr.Clear(recursive=True)) + + @staticmethod + def _create_interpolated_trajectory(plan: Any, interpolation_steps: int) -> list[torch.Tensor]: + positions = plan.position + if len(positions) < 2: + p0 = positions[0] + return [p0 if isinstance(p0, torch.Tensor) else torch.tensor(p0)] + waypoints = [p if isinstance(p, torch.Tensor) else torch.tensor(p) for p in positions] + out: list[torch.Tensor] = [] + for i in range(len(waypoints) - 1): + start, end = waypoints[i], waypoints[i + 1] + for step in range(interpolation_steps): + alpha = step / interpolation_steps + out.append(start * (1.0 - alpha) + end * alpha) + out.append(waypoints[-1]) + return out + + def mark_idle(self) -> None: + """Emit empty animation frames so stale spheres/markers don't linger between plans.""" + empty = np.empty((0, 3), dtype=float) + rr.set_time("plan", sequence=self._current_frame) + self._current_frame += 1 + rr.log("world/anim/ee", rr.Points3D(positions=empty)) + rr.set_time("sphere_animation", sequence=self._current_frame) + rr.log("world/robot_animation", rr.Points3D(positions=empty)) + rr.log("world/attached_animation", rr.Points3D(positions=empty)) diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 90c7f1c..929c99c 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -21,9 +21,8 @@ * ``mimicgen`` — single-arm MimicGen. * ``dexmimicgen`` — two-arm MimicGen with subtask coordination constraints. -* ``skillgen`` — single-arm SkillGen. SkillGen depends on a motion-planner interface; until the - planner code is ported into this repo, the CLI satisfies that interface with the upstream Arena - ``CuroboPlanner``. +* ``skillgen`` — single-arm SkillGen. SkillGen needs a motion planner; ``--planner_backend`` + selects the cuRobo backend that provides it (see :func:`_build_motion_planners`). The CLI composes a :class:`Datastream` from the task descriptor YAML, the embodiment YAML, the live env, and the HDF5 source dataset, then hands it to :class:`DataGenerator`. @@ -38,6 +37,9 @@ # Hardcoded to keep argparse importable without pulling in the heavy core package. # Add new algorithms here when registering them in isaac_autodata_core.algorithms. _ALG_CHOICES = ["mimicgen", "dexmimicgen", "skillgen"] +# Mirrors isaac_autodata_interfaces.motion_planners.PLANNER_BACKENDS; hardcoded for the same +# reason, since resolving a backend imports its cuRobo version. +_PLANNER_BACKEND_CHOICES = ["curobo", "curobo_v2"] parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument( @@ -99,6 +101,13 @@ action="store_true", help="Visualize SkillGen motion plans in a Rerun viewer (env 0 only; requires the rerun package).", ) +parser.add_argument( + "--planner_backend", + type=str, + choices=_PLANNER_BACKEND_CHOICES, + default="curobo", + help="Motion-planner backend for skillgen. Each backend needs its matching cuRobo version installed.", +) AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() @@ -229,32 +238,51 @@ def setup_async_generation( def _build_motion_planners( - datastream, num_envs: int, env_name: str, *, planner_profile: str | None = None, visualize_plan: bool = False + datastream, + num_envs: int, + env_name: str, + *, + backend: str = "curobo", + planner_profile: str | None = None, + visualize_plan: bool = False, ) -> dict: - """Construct one cuRobo v1 motion planner per env_id satisfying the SkillGen interface. + """Construct one motion planner per env_id satisfying the SkillGen interface. Planners read all world state (collision-geometry source, object poses, joint configuration) through the shared :class:`Datastream`, so they never touch the env/robot handles directly. The planner config comes from ``planner_profile`` (named by the environment profile) when given, else from task-name matching. Rerun plan visualization is opt-in via ``visualize_plan`` and limited to env 0. + + Args: + datastream: Shared read facade the planners pull world state from. + num_envs: Number of parallel environments; one planner is built per env id. + env_name: Gym task id, used for planner-config lookup when no profile is given. + backend: Planner backend name (see + :data:`isaac_autodata_interfaces.motion_planners.PLANNER_BACKENDS`). + planner_profile: Planner-profile name from the environment profile, or None. + visualize_plan: Enable the Rerun plan visualizer on env 0. + + Returns: + Mapping of env id to the planner serving it. """ - from isaac_autodata_interfaces.motion_planners.curobo.curobo_planner import CuroboPlanner - from isaac_autodata_interfaces.motion_planners.curobo.curobo_planner_cfg import CuroboPlannerCfg + from isaac_autodata_interfaces.motion_planners import get_planner_backend + + planner_cls, config_cls = get_planner_backend(backend) - planners: dict[int, CuroboPlanner] = {} + planners: dict[int, object] = {} for env_id in range(num_envs): if planner_profile is not None: - planner_config = CuroboPlannerCfg.from_profile(planner_profile) + planner_config = config_cls.from_profile(planner_profile) else: - planner_config = CuroboPlannerCfg.from_task_name(env_name) + planner_config = config_cls.from_task_name(env_name) # Visualization is rerun-based; limit to env_id 0 to keep simulation responsive. if env_id == 0: planner_config.visualize_plan = planner_config.visualize_plan or visualize_plan else: planner_config.visualize_spheres = False planner_config.visualize_plan = False - planners[env_id] = CuroboPlanner( + planners[env_id] = planner_cls( datastream=datastream, config=planner_config, env_id=env_id, @@ -330,6 +358,7 @@ def main() -> None: datastream, args_cli.num_envs, env_name, + backend=args_cli.planner_backend, planner_profile=env_profile.planner if env_profile else None, visualize_plan=args_cli.visualize_plan, )