Merge pull request #1090 from StanfordVL/fixes/robots_set_joints

make robots always call get/set_joint_positions/velocities rather tha…

omnigibson/robots/holonomic_base_robot.py

@@ -218,13 +218,14 @@ def _initialize(self):
         self.reload_controllers(self._controller_config)
 
     def apply_action(self, action):
-        j_pos = self.joints["base_footprint_rz_joint"].get_state()[0]
+        rz_joint_dof_indices = rz_joint_dof_indices = self.joints["base_footprint_rz_joint"].dof_indices
+        j_pos = self.get_joint_positions()[rz_joint_dof_indices]
         # In preparation for the base controller's @update_goal, we need to wrap the current joint pos
         # to be in range [-pi, pi], so that once the command (a delta joint pos in range [-pi, pi])
         # is applied, the final target joint pos is in range [-pi * 2, pi * 2], which is required by Isaac.
         if j_pos < -math.pi or j_pos > math.pi:
             j_pos = wrap_angle(j_pos)
-            self.joints["base_footprint_rz_joint"].set_pos(j_pos, drive=False)
+            self.set_joint_positions(j_pos, indices=rz_joint_dof_indices, drive=False)
         super().apply_action(action)
 
     @cached_property
@@ -296,7 +297,6 @@ def set_position_orientation(self, position=None, orientation=None, frame: Liter
             assert self.scene is not None, "cannot set position and orientation relative to scene without a scene"
             position, orientation = self.scene.convert_scene_relative_pose_to_world(position, orientation)
 
-        # TODO: Reconsider the need for this. Why can't these behaviors be unified? Does the joint really need to move?
         # If the simulator is playing, set the 6 base joints to achieve the desired pose of base_footprint link frame
         if og.sim.is_playing() and self.initialized:
             # Find the relative transformation from base_footprint_link ("base_footprint") frame to root_link
@@ -306,12 +306,8 @@ def set_position_orientation(self, position=None, orientation=None, frame: Liter
             inv_joint_pos, inv_joint_orn = T.invert_pose_transform(joint_pos, joint_orn)
             relative_pos, relative_orn = T.pose_transform(inv_joint_pos, inv_joint_orn, position, orientation)
             intrinsic_eulers = T.mat2euler_intrinsic(T.quat2mat(relative_orn))
-            self.joints["base_footprint_x_joint"].set_pos(relative_pos[0], drive=False)
-            self.joints["base_footprint_y_joint"].set_pos(relative_pos[1], drive=False)
-            self.joints["base_footprint_z_joint"].set_pos(relative_pos[2], drive=False)
-            self.joints["base_footprint_rx_joint"].set_pos(intrinsic_eulers[0], drive=False)
-            self.joints["base_footprint_ry_joint"].set_pos(intrinsic_eulers[1], drive=False)
-            self.joints["base_footprint_rz_joint"].set_pos(intrinsic_eulers[2], drive=False)
+            joint_positions = th.concatenate((relative_pos, intrinsic_eulers))
+            self.set_joint_positions(positions=joint_positions, indices=self.base_idx, drive=False)
 
         # Else, set the pose of the robot frame, and then move the joint frame of the world_base_joint to match it
         else:
@@ -330,9 +326,7 @@ def set_linear_velocity(self, velocity: th.Tensor):
         # such velocity), which is different from the default behavior of set_linear_velocity for all other objects.
         orn = self.root_link.get_position_orientation()[1]
         velocity_in_root_link = T.quat2mat(orn).T @ velocity
-        self.joints["base_footprint_x_joint"].set_vel(velocity_in_root_link[0], drive=False)
-        self.joints["base_footprint_y_joint"].set_vel(velocity_in_root_link[1], drive=False)
-        self.joints["base_footprint_z_joint"].set_vel(velocity_in_root_link[2], drive=False)
+        self.set_joint_velocities(velocity_in_root_link, indices=self.base_idx[:3], drive=False)
 
     def get_linear_velocity(self) -> th.Tensor:
         # Note that the link we are interested in is self.base_footprint_link, not self.root_link
@@ -353,9 +347,7 @@ def set_angular_velocity(self, velocity: th.Tensor) -> None:
         delta_intrinsic_eulers = desired_intrinsic_eulers - cur_joint_pos
         velocity_intrinsic = delta_intrinsic_eulers / delta_t
 
-        self.joints["base_footprint_rx_joint"].set_vel(velocity_intrinsic[0], drive=False)
-        self.joints["base_footprint_ry_joint"].set_vel(velocity_intrinsic[1], drive=False)
-        self.joints["base_footprint_rz_joint"].set_vel(velocity_intrinsic[2], drive=False)
+        self.set_joint_velocities(velocity_intrinsic, indices=self.base_idx[3:], drive=False)
 
     def get_angular_velocity(self) -> th.Tensor:
         # Note that the link we are interested in is self.base_footprint_link, not self.root_link

omnigibson/robots/manipulation_robot.py

@@ -175,9 +175,6 @@ def __init__(
         self._eef_to_fingertip_lengths = None  # dict mapping arm name to finger name to offset
 
         # Initialize other variables used for assistive grasping
-        self._ag_freeze_joint_pos = {
-            arm: {} for arm in self.arm_names
-        }  # Frozen positions for keeping fingers held still
         self._ag_obj_in_hand = {arm: None for arm in self.arm_names}
         self._ag_obj_constraints = {arm: None for arm in self.arm_names}
         self._ag_obj_constraint_params = {arm: {} for arm in self.arm_names}
@@ -549,11 +546,6 @@ def deploy_control(self, control, control_type):
         if self.grasping_mode != "physical" and not self._disable_grasp_handling:
             self._handle_assisted_grasping()
 
-        # Potentially freeze gripper joints
-        for arm in self.arm_names:
-            if self._ag_freeze_gripper[arm]:
-                self._freeze_gripper(arm)
-
     def _release_grasp(self, arm="default"):
         """
         Magic action to release this robot's grasp on an object
@@ -1189,20 +1181,6 @@ def _handle_release_window(self, arm="default"):
             self._ag_obj_in_hand[arm] = None
             self._ag_release_counter[arm] = None
 
-    def _freeze_gripper(self, arm="default"):
-        """
-        Freezes gripper finger joints - used in assisted grasping.
-
-        Args:
-            arm (str): specific arm to freeze gripper.
-                Default is "default" which corresponds to the first entry in self.arm_names
-        """
-        arm = self.default_arm if arm == "default" else arm
-        for joint_name, j_val in self._ag_freeze_joint_pos[arm].items():
-            joint = self._joints[joint_name]
-            joint.set_pos(pos=j_val)
-            joint.set_vel(vel=0.0)
-
     @property
     def curobo_path(self):
         """
@@ -1528,9 +1506,6 @@ def _establish_grasp_rigid(self, arm="default", ag_data=None, contact_pos=None):
         }
         self._ag_obj_in_hand[arm] = ag_obj
         self._ag_freeze_gripper[arm] = True
-        for joint in self.finger_joints[arm]:
-            j_val = joint.get_state()[0][0]
-            self._ag_freeze_joint_pos[arm][joint.joint_name] = j_val
 
     def _handle_assisted_grasping(self):
         """
@@ -1717,9 +1692,6 @@ def _establish_grasp_cloth(self, arm="default", ag_data=None):
         }
         self._ag_obj_in_hand[arm] = ag_obj
         self._ag_freeze_gripper[arm] = True
-        for joint in self.finger_joints[arm]:
-            j_val = joint.get_state()[0][0]
-            self._ag_freeze_joint_pos[arm][joint.joint_name] = j_val
 
     def _dump_state(self):
         # Call super first