Revive action primitives examples and tests

omnigibson/action_primitives/starter_semantic_action_primitives.py

@@ -99,6 +99,7 @@
 m.JOINT_POS_DIFF_THRESHOLD = 0.01
 m.JOINT_CONTROL_MIN_ACTION = 0.0
 m.MAX_ALLOWED_JOINT_ERROR_FOR_LINEAR_MOTION = math.radians(45)
+m.TIME_WITHOUT_CHECKING = 1.0
 
 log = create_module_logger(module_name=__name__)
 
@@ -116,8 +117,8 @@ def __init__(self):
         self.relative_poses = {}
         self.links_relative_poses = {}
         self.reset_pose = {
-            "original": ([0, 0, -5.0], [0, 0, 0, 1]),
-            "simplified": ([5, 0, -5.0], [0, 0, 0, 1]),
+            "original": (th.tensor([0, 0, -5.0], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32)),
+            "simplified": (th.tensor([5, 0, -5.0], dtype=th.float32), th.tensor([0, 0, 0, 1], dtype=th.float32)),
         }
 
 
@@ -164,14 +165,17 @@ def _assemble_robot_copy(self):
         if m.TIAGO_TORSO_FIXED:
             assert self.arm == "left", "Fixed torso mode only supports left arm!"
             joint_control_idx = self.robot.arm_control_idx["left"]
-            joint_pos = th.tensor(self.robot.get_joint_positions()[joint_control_idx])
+            joint_pos = th.tensor(self.robot.get_joint_positions()[joint_control_idx], dtype=th.float32)
         else:
             joint_combined_idx = th.cat([self.robot.trunk_control_idx, self.robot.arm_control_idx[fk_descriptor]])
-            joint_pos = th.tensor(self.robot.get_joint_positions()[joint_combined_idx])
+            joint_pos = th.tensor(self.robot.get_joint_positions()[joint_combined_idx], dtype=th.float32)
         link_poses = self.fk_solver.get_link_poses(joint_pos, arm_links)
 
-        # Set position of robot copy root prim
-        self._set_prim_pose(self.robot_copy.prims[self.robot_copy_type], self.robot.get_position_orientation())
+        # Set position of robot copy root prim as a tensor tuple
+        pos, orn = self.robot.get_position_orientation()
+        pos = th.tensor(pos, dtype=th.float32)
+        orn = th.tensor(orn, dtype=th.float32)
+        self._set_prim_pose(self.robot_copy.prims[self.robot_copy_type], (pos, orn))
 
         # Assemble robot meshes
         for link_name, meshes in self.robot_copy.meshes[self.robot_copy_type].items():
@@ -191,9 +195,9 @@ def _assemble_robot_copy(self):
                 self._set_prim_pose(copy_mesh, mesh_copy_pose)
 
     def _set_prim_pose(self, prim, pose):
-        translation = lazy.pxr.Gf.Vec3d(*(th.tensor(pose[0], dtype=th.float32).tolist()))
+        translation = lazy.pxr.Gf.Vec3d(*pose[0].tolist())
         prim.GetAttribute("xformOp:translate").Set(translation)
-        orientation = th.tensor(pose[1], dtype=th.float32)[[3, 0, 1, 2]]
+        orientation = pose[1][[3, 0, 1, 2]]
         prim.GetAttribute("xformOp:orient").Set(lazy.pxr.Gf.Quatd(*orientation.tolist()))
 
     def _construct_disabled_collision_pairs(self):
@@ -237,7 +241,7 @@ def _construct_disabled_collision_pairs(self):
 
         # Disable original robot colliders so copy can't collide with it
         disabled_colliders += [link.prim_path for link in self.robot.links.values()]
-        filter_categories = ["floors"]
+        filter_categories = ["floors", "carpet"]
         for obj in self.env.scene.objects:
             if obj.category in filter_categories:
                 disabled_colliders += [link.prim_path for link in obj.links.values()]
@@ -376,9 +380,9 @@ def _load_robot_copy(self):
             lazy.omni.usd.commands.CreatePrimCommand("Xform", rc["copy_path"]).do()
             copy_robot = lazy.omni.isaac.core.utils.prims.get_prim_at_path(rc["copy_path"])
             reset_pose = robot_copy.reset_pose[robot_type]
-            translation = lazy.pxr.Gf.Vec3d(*th.tensor(reset_pose[0], dtype=th.float32).tolist())
+            translation = lazy.pxr.Gf.Vec3d(*reset_pose[0].tolist())
             copy_robot.GetAttribute("xformOp:translate").Set(translation)
-            orientation = th.tensor(reset_pose[1], dtype=th.float32)[[3, 0, 1, 2]]
+            orientation = reset_pose[1][[3, 0, 1, 2]]
             copy_robot.GetAttribute("xformOp:orient").Set(lazy.pxr.Gf.Quatd(*orientation.tolist()))
 
             robot_to_copy = None
@@ -571,7 +575,6 @@ def _open_or_close(self, obj, should_open):
 
                 # If the grasp pose is too far, navigate
                 yield from self._navigate_if_needed(obj, pose_on_obj=grasp_pose)
-
                 yield from self._move_hand(grasp_pose, stop_if_stuck=True)
 
                 # We can pre-grasp in sticky grasping mode only for opening
@@ -824,14 +827,12 @@ def _place_with_predicate(self, obj, predicate):
         """
         # Update the tracking to track the object.
         self._tracking_object = obj
-
         obj_in_hand = self._get_obj_in_hand()
         if obj_in_hand is None:
             raise ActionPrimitiveError(
                 ActionPrimitiveError.Reason.PRE_CONDITION_ERROR,
                 "You need to be grasping an object first to place it somewhere.",
             )
-
         # Sample location to place object
         obj_pose = self._sample_pose_with_object_and_predicate(predicate, obj_in_hand, obj)
         hand_pose = self._get_hand_pose_for_object_pose(obj_pose)
@@ -991,17 +992,16 @@ def _move_hand_joint(self, joint_pos):
                 torso_fixed=m.TIAGO_TORSO_FIXED,
             )
 
-        # plan = self._add_linearly_interpolated_waypoints(plan, 0.1)
         if plan is None:
             raise ActionPrimitiveError(
                 ActionPrimitiveError.Reason.PLANNING_ERROR,
                 "There is no accessible path from where you are to the desired joint position. Try again",
             )
 
         # Follow the plan to navigate.
-        indented_print("Plan has %d steps", len(plan))
+        indented_print(f"Plan has {len(plan)} steps")
         for i, joint_pos in enumerate(plan):
-            indented_print("Executing grasp plan step %d/%d", i + 1, len(plan))
+            indented_print(f"Executing arm movement plan step {i + 1}/{len(plan)}")
             yield from self._move_hand_direct_joint(joint_pos, ignore_failure=True)
 
     def _move_hand_ik(self, eef_pose, stop_if_stuck=False):
@@ -1026,19 +1026,18 @@ def _move_hand_ik(self, eef_pose, stop_if_stuck=False):
                 torso_fixed=m.TIAGO_TORSO_FIXED,
             )
 
-        # plan = self._add_linearly_interpolated_waypoints(plan, 0.1)
         if plan is None:
             raise ActionPrimitiveError(
                 ActionPrimitiveError.Reason.PLANNING_ERROR,
                 "There is no accessible path from where you are to the desired joint position. Try again",
             )
 
         # Follow the plan to navigate.
-        indented_print("Plan has %d steps", len(plan))
+        indented_print(f"Plan has {len(plan)} steps")
         for i, target_pose in enumerate(plan):
             target_pos = target_pose[:3]
             target_quat = T.axisangle2quat(target_pose[3:])
-            indented_print("Executing grasp plan step %d/%d", i + 1, len(plan))
+            indented_print(f"Executing grasp plan step {i + 1}/{len(plan)}")
             yield from self._move_hand_direct_ik(
                 (target_pos, target_quat), ignore_failure=True, in_world_frame=False, stop_if_stuck=stop_if_stuck
             )
@@ -1081,22 +1080,23 @@ def _move_hand_direct_joint(self, joint_pos, stop_on_contact=False, ignore_failu
         # Store the previous eef pose for checking if we got stuck
         prev_eef_pos = th.zeros(3)
 
-        for _ in range(m.MAX_STEPS_FOR_HAND_MOVE_JOINT):
+        for i in range(m.MAX_STEPS_FOR_HAND_MOVE_JOINT):
             current_joint_pos = self.robot.get_joint_positions()[self._manipulation_control_idx]
             diff_joint_pos = th.tensor(joint_pos) - th.tensor(current_joint_pos)
             if th.max(th.abs(diff_joint_pos)).item() < m.JOINT_POS_DIFF_THRESHOLD:
                 return
             if stop_on_contact and detect_robot_collision_in_sim(self.robot, ignore_obj_in_hand=False):
                 return
-            if th.max(th.abs(self.robot.get_eef_position(self.arm) - prev_eef_pos)).item() < 0.0001:
+            # check if the eef stayed in the same pose for sufficiently long
+            if og.sim.get_sim_step_dt() * i > m.TIME_WITHOUT_CHECKING and th.max(th.abs(self.robot.get_eef_position(self.arm) - prev_eef_pos)).item() < 0.0001:
                 # We're stuck!
                 break
 
             action = self._empty_action()
             if use_delta:
-                action[self.robot.controller_action_idx[controller_name]] = diff_joint_pos
+                action[self.robot.controller_action_idx[controller_name]] = th.tensor(diff_joint_pos, dtype=th.float32)
             else:
-                action[self.robot.controller_action_idx[controller_name]] = joint_pos
+                action[self.robot.controller_action_idx[controller_name]] = th.tensor(joint_pos, dtype=th.float32)
 
             prev_eef_pos = self.robot.get_eef_position(self.arm)
             yield self._postprocess_action(action)
@@ -1205,8 +1205,10 @@ def _move_hand_linearly_cartesian(
         # into 1cm-long pieces
         start_pos, start_orn = self.robot.eef_links[self.arm].get_position_orientation()
         travel_distance = th.norm(target_pose[0] - start_pos)
-        num_poses = th.max([2, int(travel_distance / m.MAX_CARTESIAN_HAND_STEP) + 1]).item()
-        pos_waypoints = th.linspace(start_pos, target_pose[0], num_poses)
+        num_poses = int(
+            th.max(th.tensor([2, int(travel_distance / m.MAX_CARTESIAN_HAND_STEP) + 1], dtype=th.float32)).item()
+        )
+        pos_waypoints = self.linspace_1d_tensor(start_pos, target_pose[0], num_poses)
 
         # Also interpolate the rotations
         t_values = th.linspace(0, 1, num_poses)
@@ -1236,9 +1238,9 @@ def _move_hand_linearly_cartesian(
 
                 # Also decide if we can stop early.
                 current_pos, current_orn = self.robot.eef_links[self.arm].get_position_orientation()
-                pos_diff = th.norm(th.tensor(current_pos) - th.tensor(target_pose[0]))
+                pos_diff = th.norm(th.tensor(current_pos - target_pose[0], dtype=th.float32))
                 orn_diff = T.get_orientation_diff_in_radian(target_pose[1], current_orn).item()
-                if pos_diff < 0.005 and orn_diff < th.deg2rad(th.tensor([0.1])).item():
+                if pos_diff < 0.002 and orn_diff < th.deg2rad(th.tensor([0.1])).item():
                     return
 
                 if stop_on_contact and detect_robot_collision_in_sim(self.robot, ignore_obj_in_hand=False):
@@ -1432,6 +1434,10 @@ def _empty_action(self):
             action_idx = self.robot.controller_action_idx[name]
             no_op_goal = controller.compute_no_op_goal(self.robot.get_control_dict())
 
+            # TODO: wip solution for goal pose not consistent with action pose. if the controller uses delta motion, convert the no_op to all zeros
+            if self.robot._controller_config[name].get("use_delta_commands", True):
+                no_op_goal = {key: th.zeros_like(value) for key, value in no_op_goal.items()}
+
             if self.robot._controller_config[name]["name"] == "InverseKinematicsController":
                 assert (
                     self.robot._controller_config["arm_" + self.arm]["mode"] == "pose_absolute_ori"
@@ -1562,11 +1568,11 @@ def _navigate_to_pose(self, pose_2d):
                 "Could not make a navigation plan to get to the target position",
             )
 
+        # #Follow the plan to navigate.
         # self._draw_plan(plan)
-        # Follow the plan to navigate.
-        indented_print("Plan has %d steps", len(plan))
+        indented_print(f"Navigation plan has {len(plan)} steps")
         for i, pose_2d in enumerate(plan):
-            indented_print("Executing navigation plan step %d/%d", i + 1, len(plan))
+            indented_print(f"Executing navigation plan step {i + 1}/{len(plan)}")
             low_precision = True if i < len(plan) - 1 else False
             yield from self._navigate_to_pose_direct(pose_2d, low_precision=low_precision)
 
@@ -1761,7 +1767,7 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs):
                 distance_lo, distance_hi = 0.0, 5.0
                 distance = (th.rand(1) * (distance_hi - distance_lo) + distance_lo).item()
                 yaw_lo, yaw_hi = -math.pi, math.pi
-                yaw = (th.rand(1) * (yaw_hi - yaw_lo) + yaw_lo).item()
+                yaw = th.tensor((th.rand(1) * (yaw_hi - yaw_lo) + yaw_lo).item(), dtype=th.float32)
                 avg_arm_workspace_range = th.mean(self.robot.arm_workspace_range[self.arm])
                 pose_2d = th.tensor(
                     [
@@ -1783,6 +1789,7 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs):
                 if not self._test_pose(pose_2d, context, pose_on_obj=pose_on_obj, **kwargs):
                     continue
 
+                indented_print("Found valid position near object.")
                 return pose_2d
 
             raise ActionPrimitiveError(
@@ -1871,7 +1878,7 @@ def _sample_pose_with_object_and_predicate(
 
             # Get the object pose by subtracting the offset
             sampled_obj_pose = T.pose2mat((sampled_bb_center, sampled_bb_orn)) @ T.pose_inv(
-                T.pose2mat((bb_center_in_base, [0, 0, 0, 1]))
+                T.pose2mat((bb_center_in_base, th.tensor([0, 0, 0, 1], dtype=th.float32)))
             )
 
             # Check that the pose is near one of the poses in the near_poses list if provided.
@@ -1923,11 +1930,10 @@ def _get_robot_pose_from_2d_pose(pose_2d):
             pose_2d (Iterable): (x, y, yaw) 2d pose
 
         Returns:
-            2-tuple:
-                - 3-array: (x,y,z) Position in the world frame
-                - 4-array: (x,y,z,w) Quaternion orientation in the world frame
+            th.tensor: (x,y,z) Position in the world frame
+            th.tensor: (x,y,z,w) Quaternion orientation in the world frame
         """
-        pos = th.tensor([pose_2d[0], pose_2d[1], m.DEFAULT_BODY_OFFSET_FROM_FLOOR])
+        pos = th.tensor([pose_2d[0], pose_2d[1], m.DEFAULT_BODY_OFFSET_FROM_FLOOR], dtype=th.float32)
         orn = T.euler2quat(th.tensor([0, 0, pose_2d[2]], dtype=th.float32))
         return pos, orn
 
@@ -1991,3 +1997,24 @@ def _settle_robot(self):
                 break
             empty_action = self._empty_action()
             yield self._postprocess_action(empty_action)
+
+    def linspace_1d_tensor(self, start_pos, target_pose, num_poses):
+        """
+        Create evenly spaced samples between two 1D tensors.
+
+        :param start_pos: Starting 1D tensor
+        :param target_pose: Ending 1D tensor
+        :param num_poses: Number of poses (samples) to generate
+        :return: Tensor of shape (num_poses, dim) where dim is the dimension of input tensors
+        """
+        # Ensure inputs are 1D tensors
+        assert start_pos.dim() == 1 and target_pose.dim() == 1, "Input tensors must be 1D"
+        assert start_pos.shape == target_pose.shape, "Input tensors must have the same shape"
+
+        # Create a tensor of interpolation factors
+        t = th.linspace(0, 1, num_poses)
+
+        # Perform the interpolation
+        interpolated_points = start_pos.unsqueeze(0) + (target_pose - start_pos).unsqueeze(0) * t.unsqueeze(1)
+
+        return interpolated_points