`motion_gen.warmup` resulting in breaking cuda graph when using `link_poses` with `plan_batch`

[tylerlum]

Problem

This is the exact same code as the Motion Generation example code but with 2 changes:

• I do plan_batch and repeat the inputs 2 times
• I add link_poses=state.link_pose to the plan_batch function.

import torch

# cuRobo
from curobo.types.math import Pose
from curobo.types.robot import JointState
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig

world_config = {
    "mesh": {
        "base_scene": {
            "pose": [10.5, 0.080, 1.6, 0.043, -0.471, 0.284, 0.834],
            "file_path": "scene/nvblox/srl_ur10_bins.obj",
        },
    },
    "cuboid": {
        "table": {
            "dims": [5.0, 5.0, 0.2],  # x, y, z
            "pose": [0.0, 0.0, -0.1, 1, 0, 0, 0.0],  # x, y, z, qw, qx, qy, qz
        },
    },
}

motion_gen_config = MotionGenConfig.load_from_robot_config(
    "ur5e.yml",
    world_config,
    interpolation_dt=0.01,
)
motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup()

retract_cfg = motion_gen.get_retract_config()

### CHANGE START ###
# state = motion_gen.rollout_fn.compute_kinematics(
#     JointState.from_position(retract_cfg.view(1, -1))
# )
state = motion_gen.rollout_fn.compute_kinematics(
    JointState.from_position(retract_cfg.view(1, -1)).repeat_seeds(2)
)

# goal_pose = Pose.from_list([-0.4, 0.0, 0.4, 1.0, 0.0, 0.0, 0.0])  # x, y, z, qw, qx, qy, qz
goal_pose = Pose.from_list([-0.4, 0.0, 0.4, 1.0, 0.0, 0.0, 0.0]).repeat(2)  # x, y, z, qw, qx, qy, qz
# start_state = JointState.from_position(
#     torch.zeros(1, 6).cuda(),
#     joint_names=[
#         "shoulder_pan_joint",
#         "shoulder_lift_joint",
#         "elbow_joint",
#         "wrist_1_joint",
#         "wrist_2_joint",
#         "wrist_3_joint",
#     ],
# )
start_state = JointState.from_position(
    torch.zeros(1, 6).cuda(),
    joint_names=[
        "shoulder_pan_joint",
        "shoulder_lift_joint",
        "elbow_joint",
        "wrist_1_joint",
        "wrist_2_joint",
        "wrist_3_joint",
    ],
).repeat_seeds(2)

# result = motion_gen.plan_single(start_state, goal_pose, MotionGenPlanConfig(max_attempts=1))
result = motion_gen.plan_batch(start_state, goal_pose, MotionGenPlanConfig(max_attempts=1), link_poses=state.link_pose)
# traj = result.get_interpolated_plan()  # result.optimized_dt has the dt between timesteps
### CHANGE END ###

print("Trajectory Generated: ", result.success)

I get the following error with this code:

python test.py
Creating new Mesh cache: 1
changing goal type, breaking previous cuda graph.
NoneType: None
Traceback (most recent call last):
  File "test.py", line 67, in <module>
    result = motion_gen.plan_batch(start_state, goal_pose, MotionGenPlanConfig(max_attempts=1), link_poses=state.link_pose)
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/motion_gen.py", line 1533, in plan_batch
    result = self._plan_batch_attempts(
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/motion_gen.py", line 2933, in _plan_batch_attempts
    result = self._plan_from_solve_state_batch(
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/motion_gen.py", line 3570, in _plan_from_solve_state_batch
    ik_result = self._solve_ik_from_solve_state(
  File "/juno/u/tylerlum/miniconda3/envs/nerf_grasping_env_experimental/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/motion_gen.py", line 2571, in _solve_ik_from_solve_state
    ik_result = self.ik_solver.solve_any(
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/ik_solver.py", line 1182, in solve_any
    return self.solve_batch(
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/ik_solver.py", line 794, in solve_batch
    return self._solve_from_solve_state(
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/ik_solver.py", line 1037, in _solve_from_solve_state
    goal_buffer = self._update_goal_buffer(solve_state, goal_pose, retract_config, link_poses)
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/wrap/reacher/ik_solver.py", line 586, in _update_goal_buffer
    log_error("changing goal type, breaking previous cuda graph.")
  File "/juno/u/tylerlum/github_repos/curobo/src/curobo/util/logger.py", line 103, in log_error
    raise ValueError(txt)
ValueError: changing goal type, breaking previous cuda graph.

Root Cause

I looked into this and the reason it happens is that the warmup step assumes link_poses=None when doing plan_batch, so it breaks the cuda graph when it is not None afterwards

1. cuRobo installation mode (choose from [python, isaac sim, docker python, docker isaac sim]): python
2. python version: Python 3.8.19
3. Isaac Sim version (if using): Not Used

---

Issue Details

[balakumar-s]

Added link_poses to warmup.