Integrate cuRobo into action primitives & primitives refactoring

omnigibson/action_primitives/action_primitive_set_base.py

@@ -58,14 +58,9 @@ def __init_subclass__(cls, **kwargs):
         if not inspect.isabstract(cls):
             REGISTERED_PRIMITIVE_SETS[cls.__name__] = cls
 
-    def __init__(self, env):
+    def __init__(self, env, robot):
         self.env = env
-
-    @property
-    def robot(self):
-        # Currently returns the first robot in the environment, but can be scaled to multiple robots
-        # by creating multiple action generators and passing in a robot index etc.
-        return self.env.robots[0]
+        self.robot = robot
 
     @abstractmethod
     def get_action_space(self):
@@ -84,7 +79,7 @@ def apply(self, action):
         pass
 
     @abstractmethod
-    def apply_ref(self, action, *args):
+    def apply_ref(self, primitive, *args):
         """
         Apply a primitive action by reference.
 

omnigibson/action_primitives/curobo.py

@@ -228,7 +228,7 @@ def __init__(
                 num_trajopt_seeds=4,
                 num_graph_seeds=4,
                 interpolation_dt=0.03,
-                collision_activation_distance=0.005,
+                collision_activation_distance=0.05,
                 self_collision_check=True,
                 maximum_trajectory_dt=None,
                 fixed_iters_trajopt=True,
@@ -336,6 +336,7 @@ def check_collisions(
         self,
         q,
         check_self_collision=True,
+        skip_obstacle_update=False,
     ):
         """
         Checks collisions between the sphere representation of the robot and the rest of the current scene
@@ -353,7 +354,8 @@ def check_collisions(
         emb_sel = CuroboEmbodimentSelection.DEFAULT
 
         # Update obstacles
-        self.update_obstacles()
+        if not skip_obstacle_update:
+            self.update_obstacles()
 
         q_pos = self.robot.get_joint_positions().unsqueeze(0)
         cu_joint_state = lazy.curobo.types.state.JointState(
@@ -428,7 +430,9 @@ def compute_trajectories(
         return_full_result=False,
         success_ratio=None,
         attached_obj=None,
+        motion_constraint=None,
         attached_obj_scale=None,
+        skip_obstacle_update=False,
         emb_sel=CuroboEmbodimentSelection.DEFAULT,
     ):
         """
@@ -509,8 +513,16 @@ def compute_trajectories(
             success_ratio=1.0 / self.batch_size if success_ratio is None else success_ratio,
         )
 
-        # Refresh the collision state
-        self.update_obstacles()
+        # Add the pose cost metric
+        if motion_constraint is None:
+            motion_constraint = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+        pose_cost_metric = lazy.curobo.wrap.reacher.motion_gen.PoseCostMetric(
+            hold_partial_pose=False, hold_vec_weight=self._tensor_args.to_device(motion_constraint)
+        )
+        plan_cfg.pose_cost_metric = pose_cost_metric
+
+        if not skip_obstacle_update:
+            self.update_obstacles()
 
         for link_name in target_pos.keys():
             target_pos_link = target_pos[link_name]
@@ -581,9 +593,12 @@ def compute_trajectories(
                     object_names=obj_paths,
                     ee_pose=ee_pose,
                     link_name=self.robot.curobo_attached_object_link_names[ee_link_name],
-                    scale=0.99 if attached_obj_scale is None else attached_obj_scale[ee_link_name],
+                    scale=1.0 if attached_obj_scale is None else attached_obj_scale[ee_link_name],
                     pitch_scale=1.0,
                     merge_meshes=True,
+                    world_objects_pose_offset=lazy.curobo.types.math.Pose.from_list(
+                        [0, 0, 0.01, 1, 0, 0, 0], self._tensor_args
+                    ),
                 )
 
         all_rollout_fns = [