Merge pull request #1042 from StanfordVL/fix/flaky-tests

fix flaky tests

omnigibson/action_primitives/curobo.py

@@ -28,7 +28,7 @@
 m.HOLONOMIC_BASE_REVOLUTE_JOINT_LIMIT = math.pi * 2  # radians
 
 
-class CuroboEmbodimentSelection(str, Enum):
+class CuRoboEmbodimentSelection(str, Enum):
     BASE = "base"
     ARM = "arm"
     DEFAULT = "default"
@@ -176,16 +176,15 @@ def __init__(
         self.debug = debug
         self.robot = robot
         self.robot_joint_names = list(robot.joints.keys())
-        self._fk = FKSolver(self.robot.robot_arm_descriptor_yamls[robot.default_arm], self.robot.urdf_path)
         self.batch_size = batch_size
 
         # Load robot config and usd paths and make sure paths point correctly
         robot_cfg_path_dict = robot.curobo_path if robot_cfg_path is None else robot_cfg_path
         if not isinstance(robot_cfg_path_dict, dict):
-            robot_cfg_path_dict = {CuroboEmbodimentSelection.DEFAULT: robot_cfg_path_dict}
+            robot_cfg_path_dict = {CuRoboEmbodimentSelection.DEFAULT: robot_cfg_path_dict}
         if use_default_embodiment_only:
             robot_cfg_path_dict = {
-                CuroboEmbodimentSelection.DEFAULT: robot_cfg_path_dict[CuroboEmbodimentSelection.DEFAULT]
+                CuRoboEmbodimentSelection.DEFAULT: robot_cfg_path_dict[CuRoboEmbodimentSelection.DEFAULT]
             }
         robot_usd_path = robot.usd_path if robot_usd_path is None else robot_usd_path
 
@@ -265,7 +264,7 @@ def update_joint_limits(self, robot_cfg_obj, emb_sel):
 
                 joint_limits.position[1][joint_idx] = -joint_limits.position[0][joint_idx]
 
-    def save_visualization(self, q, file_path, emb_sel=CuroboEmbodimentSelection.DEFAULT):
+    def save_visualization(self, q, file_path, emb_sel=CuRoboEmbodimentSelection.DEFAULT):
         # Update obstacles
         self.update_obstacles()
 
@@ -311,12 +310,12 @@ def update_obstacles(self, ignore_paths=None):
             ],
         )
         # All embodiment selections share the same world collision checker
-        self.mg[CuroboEmbodimentSelection.DEFAULT].update_world(obstacles)
+        self.mg[CuRoboEmbodimentSelection.DEFAULT].update_world(obstacles)
         print("Synced CuRobo world from stage.")
 
     def update_obstacles_fast(self):
         # All embodiment selections share the same world collision checker
-        world_coll_checker = self.mg[CuroboEmbodimentSelection.DEFAULT].world_coll_checker
+        world_coll_checker = self.mg[CuRoboEmbodimentSelection.DEFAULT].world_coll_checker
         for i, prim_path in enumerate(world_coll_checker._env_mesh_names[0]):
             if prim_path is None:
                 continue
@@ -344,13 +343,13 @@ def check_collisions(
             q (th.tensor): (N, D)-shaped tensor, representing N-total different joint configurations to check
                 collisions against the world
             check_self_collision (bool): Whether to check self-collisions or not
-            emb_sel (CuroboEmbodimentSelection): Which embodiment selection to use for checking collisions
+            emb_sel (CuRoboEmbodimentSelection): Which embodiment selection to use for checking collisions
 
         Returns:
             th.tensor: (N,)-shaped tensor, where each value is True if in collision, else False
         """
         # check_collisions only makes sense for the default embodiment where all the joints are actuated
-        emb_sel = CuroboEmbodimentSelection.DEFAULT
+        emb_sel = CuRoboEmbodimentSelection.DEFAULT
 
         # Update obstacles
         self.update_obstacles()
@@ -396,7 +395,7 @@ def update_locked_joints(self, cu_joint_state, emb_sel):
 
         Args:
             cu_joint_state (JointState): JointState object representing the current joint positions
-            emb_sel (CuroboEmbodimentSelection): Which embodiment selection to use for updating locked joints
+            emb_sel (CuRoboEmbodimentSelection): Which embodiment selection to use for updating locked joints
         """
         kc = self.mg[emb_sel].kinematics.kinematics_config
         # Update the lock joint state position
@@ -429,7 +428,7 @@ def compute_trajectories(
         success_ratio=None,
         attached_obj=None,
         attached_obj_scale=None,
-        emb_sel=CuroboEmbodimentSelection.DEFAULT,
+        emb_sel=CuRoboEmbodimentSelection.DEFAULT,
     ):
         """
         Computes the robot joint trajectory to reach the desired @target_pos and @target_quat
@@ -462,7 +461,7 @@ def compute_trajectories(
                 link names and the values are the corresponding BaseObject instances to attach to that link
             attached_obj_scale (None or Dict[str, float]): If specified, a dictionary where the keys are the end-effector
                 link names and the values are the corresponding scale to apply to the attached object
-            emb_sel (CuroboEmbodimentSelection): Which embodiment selection to use for computing trajectories
+            emb_sel (CuRoboEmbodimentSelection): Which embodiment selection to use for computing trajectories
         Returns:
             2-tuple or list of MotionGenResult: If @return_full_result is True, will return a list of raw MotionGenResult
                 object(s) computed from internal batch trajectory computations. If it is False, will return 2-tuple
@@ -658,6 +657,7 @@ def compute_trajectories(
                 ik_goal_batch = lazy.curobo.types.math.Pose(
                     position=batch_target_pos,
                     quaternion=batch_target_quat,
+                    name=link_name,
                 )
 
                 ik_goal_batch_by_link[link_name] = ik_goal_batch
@@ -702,13 +702,13 @@ def compute_trajectories(
         else:
             return successes, paths
 
-    def path_to_joint_trajectory(self, path, emb_sel=CuroboEmbodimentSelection.DEFAULT):
+    def path_to_joint_trajectory(self, path, emb_sel=CuRoboEmbodimentSelection.DEFAULT):
         """
         Converts raw path from motion generator into joint trajectory sequence
 
         Args:
             path (JointState): Joint state path to convert into joint trajectory
-            emb_sel (CuroboEmbodimentSelection): Which embodiment to use for the robot
+            emb_sel (CuRoboEmbodimentSelection): Which embodiment to use for the robot
 
         Returns:
             torch.tensor: (T, D) tensor representing the interpolated joint trajectory
@@ -718,39 +718,59 @@ def path_to_joint_trajectory(self, path, emb_sel=CuroboEmbodimentSelection.DEFAU
         cmd_plan = self.mg[emb_sel].get_full_js(path)
         return cmd_plan.get_ordered_joint_state(self.robot_joint_names).position
 
-    def convert_q_to_eef_traj(self, traj, return_axisangle=False, emb_sel=CuroboEmbodimentSelection.DEFAULT):
+    def path_to_eef_trajectory(self, path, return_axisangle=False, emb_sel=CuRoboEmbodimentSelection.DEFAULT):
         """
-        Converts a joint trajectory @traj into an equivalent trajectory defined by end effector poses
+        Converts raw path from motion generator into end-effector trajectory sequence in the robot frame.
+        This trajectory sequence can be executed by an IKController, although there is no guaranteee that
+        the controller will output the same joint trajectory as the one computed by cuRobo.
 
         Args:
-            traj (torch.Tensor): (T, D)-shaped joint trajectory
+            path (JointState): Joint state path to convert into joint trajectory
             return_axisangle (bool): Whether to return the interpolated orientations in quaternion or axis-angle representation
-            emb_sel (CuroboEmbodimentSelection): Which embodiment to use for the robot
+            emb_sel (CuRoboEmbodimentSelection): Which embodiment to use for the robot
 
         Returns:
-            torch.Tensor: (T, [6, 7])-shaped array where each entry is is the (x,y,z) position and (x,y,z,w) quaternion
-                (if @return_axisangle is False) or (ax, ay, az) axis-angle orientation, specified in the robot frame.
+            Dict[str, torch.Tensor]: Mapping eef link names to (T, [6, 7])-shaped array where each entry is is the (x,y,z) position
+            and (x,y,z,w) quaternion (if @return_axisangle is False) or (ax, ay, az) axis-angle orientation, specified in the robot frame.
         """
-        # Prune the relevant joints from the trajectory
-        traj = traj[:, self.robot.arm_control_idx[self.robot.default_arm]]
-        n = len(traj)
-
-        # Use forward kinematic solver to compute the EEF link positions
-        positions = self._tensor_args.to_device(th.zeros((n, 3)))
-        orientations = self._tensor_args.to_device(
-            th.zeros((n, 4))
-        )  # This will be quat initially but we may convert to aa representation
-
-        for i, qpos in enumerate(traj):
-            pose = self._fk.get_link_poses(joint_positions=qpos, link_names=[self.ee_link[emb_sel]])
-            positions[i] = pose[self.ee_link[emb_sel]][0]
-            orientations[i] = pose[self.ee_link[emb_sel]][1]
-
-        # Possibly convert orientations to aa-representation
-        if return_axisangle:
-            orientations = T.quat2axisangle(orientations)
-
-        return th.concatenate([positions, orientations], dim=-1)
+        # If the base-only embodiment is selected, the eef links stay the same, return the current eef poses in the robot frame
+        if emb_sel == CuRoboEmbodimentSelection.BASE:
+            link_poses = dict()
+            for arm_name in self.robot.arm_names:
+                link_name = self.robot.eef_link_names[arm_name]
+                position, orientation = self.robot.get_relative_eef_pose(arm_name)
+                if return_axisangle:
+                    orientation = T.quat2axisangle(orientation)
+                link_poses[link_name] = th.cat([position, orientation], dim=-1)
+            return link_poses
+
+        cmd_plan = self.mg[emb_sel].get_full_js(path)
+        robot_state = self.mg[emb_sel].kinematics.compute_kinematics(path)
+
+        link_poses = dict()
+
+        for link_name, poses in robot_state.link_poses.items():
+            position = poses.position
+            # wxyz -> xyzw
+            orientation = poses.quaternion[:, [1, 2, 3, 0]]
+
+            # If the robot is holonomic, we need to transform the poses to the base link frame
+            if isinstance(self.robot, HolonomicBaseRobot):
+                base_link_position = th.zeros_like(position)
+                base_link_position[:, 0] = cmd_plan.position[:, cmd_plan.joint_names.index("base_footprint_x_joint")]
+                base_link_position[:, 1] = cmd_plan.position[:, cmd_plan.joint_names.index("base_footprint_y_joint")]
+                base_link_euler = th.zeros_like(position)
+                base_link_euler[:, 2] = cmd_plan.position[:, cmd_plan.joint_names.index("base_footprint_rz_joint")]
+                base_link_orientation = T.euler2quat(base_link_euler)
+                position, orientation = T.relative_pose_transform(
+                    position, orientation, base_link_position, base_link_orientation
+                )
+
+            if return_axisangle:
+                orientation = T.quat2axisangle(orientation)
+            link_poses[link_name] = th.cat([position, orientation], dim=-1)
+
+        return link_poses
 
     @property
     def tensor_args(self):

omnigibson/envs/data_wrapper.py

@@ -309,18 +309,19 @@ def _optimize_sim_for_data_collection(self, viewport_camera_path):
         # We have to do a super hacky workaround to avoid the GUI freezing, which is
         # toggling these settings to be True -> False -> True
         # Only setting it to True once will actually freeze the GUI for some reason!
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", True)
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", True)
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", False)
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", False)
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", True)
-        lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", True)
-
-        # Disable mouse grabbing since we're only using the UI passively
-        lazy.carb.settings.get_settings().set_bool("/physics/mouseInteractionEnabled", False)
-        lazy.carb.settings.get_settings().set_bool("/physics/mouseGrab", False)
-        lazy.carb.settings.get_settings().set_bool("/physics/forceGrab", False)
-        lazy.carb.settings.get_settings().set_bool("/physics/suppressReadback", True)
+        if not gm.HEADLESS:
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", True)
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", True)
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", False)
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", False)
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRendering", True)
+            lazy.carb.settings.get_settings().set_bool("/app/asyncRenderingLowLatency", True)
+
+            # Disable mouse grabbing since we're only using the UI passively
+            lazy.carb.settings.get_settings().set_bool("/physics/mouseInteractionEnabled", False)
+            lazy.carb.settings.get_settings().set_bool("/physics/mouseGrab", False)
+            lazy.carb.settings.get_settings().set_bool("/physics/forceGrab", False)
+            lazy.carb.settings.get_settings().set_bool("/physics/suppressReadback", True)
 
         # Set the dump filter for better performance
         # TODO: Possibly remove this feature once we have fully tensorized state saving, which may be more efficient

omnigibson/examples/robots/curobo_example.py

@@ -5,14 +5,14 @@
 import omnigibson as og
 import omnigibson.lazy as lazy
 import omnigibson.utils.transform_utils as T
-from omnigibson.action_primitives.curobo import CuroboEmbodimentSelection, CuRoboMotionGenerator
+from omnigibson.action_primitives.curobo import CuRoboEmbodimentSelection, CuRoboMotionGenerator
 from omnigibson.macros import gm, macros
 from omnigibson.object_states import Touching
 from omnigibson.utils.ui_utils import choose_from_options
 
 
 def plan_trajectory(
-    cmg, target_pos, target_quat, emb_sel=CuroboEmbodimentSelection.DEFAULT, attached_obj=None, attached_obj_scale=None
+    cmg, target_pos, target_quat, emb_sel=CuRoboEmbodimentSelection.DEFAULT, attached_obj=None, attached_obj_scale=None
 ):
     # Generate collision-free trajectories to the sampled eef poses (including self-collisions)
     successes, traj_paths = cmg.compute_trajectories(
@@ -229,21 +229,6 @@ def test_curobo():
     robot = env.robots[0]
     eef_markers = [env.scene.object_registry("name", f"eef_marker_{i}") for i in range(2)]
 
-    # Stablize the robot and update the initial state
-    robot.reset()
-
-    # Open the gripper(s) to match cuRobo's default state
-    for arm_name in robot.gripper_control_idx.keys():
-        gripper_control_idx = robot.gripper_control_idx[arm_name]
-        robot.set_joint_positions(th.ones_like(gripper_control_idx), indices=gripper_control_idx, normalized=True)
-    robot.keep_still()
-
-    for _ in range(5):
-        og.sim.step()
-
-    env.scene.update_initial_state()
-    env.scene.reset()
-
     # Create CuRobo instance
     cmg = CuRoboMotionGenerator(
         robot=robot,
@@ -254,8 +239,8 @@ def test_curobo():
     table = env.scene.object_registry("name", "table")
     cologne = env.scene.object_registry("name", "cologne")
     fridge = env.scene.object_registry("name", "fridge")
-    cologne_scale = 0.99
-    fridge_door_scale = 0.7
+    cologne_scale = 0.8
+    fridge_door_scale = 0.8
     if robot_type == "R1":
         table_local_pose = (th.tensor([-0.8, 0.0, -0.402]), th.tensor([0.0, 0.0, 0.0, 1.0]))
         cologne_local_pose = (th.tensor([-0.03, 0.0, 0.052]), T.euler2quat(th.tensor([math.pi, -math.pi / 2.0, 0.0])))
@@ -292,7 +277,7 @@ def test_curobo():
     table_nav_pos, table_nav_quat = T.pose_transform(*table.get_position_orientation(), *table_local_pose)
     target_pos = {robot.base_footprint_link_name: table_nav_pos}
     target_quat = {robot.base_footprint_link_name: table_nav_quat}
-    emb_sel = CuroboEmbodimentSelection.BASE
+    emb_sel = CuRoboEmbodimentSelection.BASE
     attached_obj = None
     attached_obj_scale = None
     plan_and_execute_trajectory(
@@ -308,7 +293,7 @@ def test_curobo():
     right_hand_pos, right_hand_quat = robot.get_eef_pose(arm="right")
     target_pos = {robot.eef_link_names["left"]: left_hand_pos}
     target_quat = {robot.eef_link_names["left"]: left_hand_quat}
-    emb_sel = CuroboEmbodimentSelection.ARM
+    emb_sel = CuRoboEmbodimentSelection.ARM
     attached_obj = None
     attached_obj_scale = None
     plan_and_execute_trajectory(
@@ -327,7 +312,7 @@ def test_curobo():
         robot.eef_link_names["left"]: left_hand_reset_quat,
         robot.eef_link_names["right"]: right_hand_reset_quat,
     }
-    emb_sel = CuroboEmbodimentSelection.ARM
+    emb_sel = CuRoboEmbodimentSelection.ARM
     attached_obj = {robot.eef_link_names["left"]: cologne.root_link}
     attached_obj_scale = {robot.eef_link_names["left"]: cologne_scale}
     plan_and_execute_trajectory(
@@ -338,7 +323,7 @@ def test_curobo():
     fridge_nav_pos, fridge_nav_quat = T.pose_transform(*fridge.get_position_orientation(), *fridge_local_pose)
     target_pos = {robot.base_footprint_link_name: fridge_nav_pos}
     target_quat = {robot.base_footprint_link_name: fridge_nav_quat}
-    emb_sel = CuroboEmbodimentSelection.BASE
+    emb_sel = CuRoboEmbodimentSelection.BASE
     attached_obj = {robot.eef_link_names["left"]: cologne.root_link}
     attached_obj_scale = {robot.eef_link_names["left"]: cologne_scale}
     plan_and_execute_trajectory(
@@ -349,7 +334,7 @@ def test_curobo():
     right_hand_pos, right_hand_quat = T.pose_transform(*fridge.get_position_orientation(), *fridge_door_local_pose)
     target_pos = {robot.eef_link_names["right"]: right_hand_pos}
     target_quat = {robot.eef_link_names["right"]: right_hand_quat}
-    emb_sel = CuroboEmbodimentSelection.ARM
+    emb_sel = CuRoboEmbodimentSelection.ARM
     attached_obj = {robot.eef_link_names["left"]: cologne.root_link}
     attached_obj_scale = {robot.eef_link_names["left"]: cologne_scale}
     plan_and_execute_trajectory(
@@ -366,7 +351,7 @@ def test_curobo():
     right_hand_pos, right_hand_quat = T.pose_transform(*fridge.get_position_orientation(), *fridge_door_open_local_pose)
     target_pos = {robot.eef_link_names["right"]: right_hand_pos}
     target_quat = {robot.eef_link_names["right"]: right_hand_quat}
-    emb_sel = CuroboEmbodimentSelection.DEFAULT
+    emb_sel = CuRoboEmbodimentSelection.DEFAULT
     attached_obj = {
         robot.eef_link_names["left"]: cologne.root_link,
         robot.eef_link_names["right"]: fridge.links["link_0"],
@@ -383,7 +368,7 @@ def test_curobo():
     left_hand_pos, left_hand_quat = T.pose_transform(*fridge.get_position_orientation(), *fridge_place_local_pose)
     target_pos = {robot.eef_link_names["left"]: left_hand_pos}
     target_quat = {robot.eef_link_names["left"]: left_hand_quat}
-    emb_sel = CuroboEmbodimentSelection.DEFAULT
+    emb_sel = CuRoboEmbodimentSelection.DEFAULT
     attached_obj = {robot.eef_link_names["left"]: cologne.root_link}
     attached_obj_scale = {robot.eef_link_names["left"]: cologne_scale}
     plan_and_execute_trajectory(