Add support for older warp versions (<1.0.0)

src/curobo/rollout/cost/bound_cost.py

@@ -1550,7 +1550,10 @@ def make_bound_pos_smooth_kernel(dof_template: int):
                 out_grad_a[b_addrs + i] = g_a[i]
                 out_grad_j[b_addrs + i] = g_j[i]
 
-    return wp.Kernel(forward_bound_smooth_loop_warp)
+    module = wp.get_module(forward_bound_smooth_loop_warp.__module__)
+    key = "forward_bound_smooth_loop_warp_" + str(dof_template)
+
+    return wp.Kernel(forward_bound_smooth_loop_warp, key=key, module=module)
 
 
 @get_cache_fn_decorator()
@@ -1647,4 +1650,6 @@ def make_bound_pos_kernel(dof_template: int):
             for i in range(dof_template):
                 out_grad_p[b_addrs + i] = g_p[i]
 
-    return wp.Kernel(forward_bound_pos_loop_warp)
+    module = wp.get_module(forward_bound_pos_loop_warp.__module__)
+    key = "forward_bound_pos_loop_warp_" + str(dof_template)
+    return wp.Kernel(forward_bound_pos_loop_warp, key=key, module=module)

src/curobo/rollout/cost/dist_cost.py

@@ -224,7 +224,9 @@ def make_l2_kernel(dof_template: int):
             for i in range(dof_template):
                 out_grad_p[b_addrs + i] = g_p[i]
 
-    return wp.Kernel(forward_l2_loop_warp)
+    module = wp.get_module(forward_l2_loop_warp.__module__)
+    key = "forward_l2_loop" + str(dof_template)
+    return wp.Kernel(forward_l2_loop_warp, key=key, module=module)
 
 
 # create a bound cost tensor:

src/curobo/rollout/cost/pose_cost.py

@@ -466,6 +466,8 @@ class PoseCost(CostBase, PoseCostConfig):
         batch_pose_idx: torch.Tensor,
         mode: PoseErrorType = PoseErrorType.BATCH_GOAL,
     ):
+        if len(query_pose.position.shape) == 2:
+            log_error("Query pose should be [batch, horizon, -1]")
         ee_goal_pos = goal_pose.position
         ee_goal_quat = goal_pose.quaternion
         self.cost_type = mode
@@ -476,9 +478,9 @@ class PoseCost(CostBase, PoseCostConfig):
         num_goals = 1
 
         distance = PoseError.apply(
-            query_pose.position.unsqueeze(1),
+            query_pose.position,
             ee_goal_pos,
-            query_pose.quaternion.unsqueeze(1),
+            query_pose.quaternion,
             ee_goal_quat,
             self.vec_weight,
             self.weight,