Merge pull request #1075 from StanfordVL/feat/robot-fixes

Feat/robot fixes

docs/tutorials/custom_robot_import.md

@@ -494,28 +494,6 @@ Now that we have the USD file for the robot, let's write our own robot class. Fo
         def wheel_axle_length(self):
             return 0.330
 
-        @property
-        def finger_lengths(self):
-            return {self.default_arm: 0.04}
-
-        @property
-        def assisted_grasp_start_points(self):
-            return {
-                self.default_arm: [
-                    GraspingPoint(link_name="link_gripper_finger_right", position=th.tensor([0.013, 0.0, 0.01])),
-                    GraspingPoint(link_name="link_gripper_finger_right", position=th.tensor([-0.01, 0.0, 0.009])),
-                ]
-            }
-
-        @property
-        def assisted_grasp_end_points(self):
-            return {
-                self.default_arm: [
-                    GraspingPoint(link_name="link_gripper_finger_left", position=th.tensor([0.013, 0.0, 0.01])),
-                    GraspingPoint(link_name="link_gripper_finger_left", position=th.tensor([-0.01, 0.0, 0.009])),
-                ]
-            }
-
         @property
         def disabled_collision_pairs(self):
             return [

docs/tutorials/demo_collection.md

@@ -101,7 +101,7 @@ playback_env = DataPlaybackWrapper.create_from_hdf5(
 )
 
 # Playback the entire dataset and record observations
-playback_env.playback_dataset(record=True)
+playback_env.playback_dataset(record_data=True)
 
 # Save the recorded playback data
 playback_env.save_data()

omnigibson/action_primitives/starter_semantic_action_primitives.py

@@ -102,7 +102,7 @@
 m.LOW_PRECISION_DIST_THRESHOLD = 0.1
 m.LOW_PRECISION_ANGLE_THRESHOLD = 0.2
 
-m.JOINT_POS_DIFF_THRESHOLD = 0.01
+m.JOINT_POS_DIFF_THRESHOLD = 0.005
 m.LOW_PRECISION_JOINT_POS_DIFF_THRESHOLD = 0.05
 m.JOINT_CONTROL_MIN_ACTION = 0.0
 m.MAX_ALLOWED_JOINT_ERROR_FOR_LINEAR_MOTION = math.radians(45)
@@ -506,7 +506,10 @@ def _sample_grasp_pose(self, obj):
         # Identity quaternion for top-down grasping (x-forward, y-right, z-down)
         approach_dir = T.quat2mat(grasp_quat) @ th.tensor([0.0, 0.0, -1.0])
 
-        pregrasp_offset = self.robot.finger_lengths[self.arm] / 2.0 + m.GRASP_APPROACH_DISTANCE
+        avg_finger_offset = th.mean(
+            th.tensor([length for length in self.robot.eef_to_fingertip_lengths[self.arm].values()])
+        )
+        pregrasp_offset = avg_finger_offset + m.GRASP_APPROACH_DISTANCE
 
         pregrasp_pos = grasp_pos - approach_dir * pregrasp_offset
 
@@ -581,7 +584,7 @@ def _grasp(self, obj):
             # By default, it's NOT the z-axis of the world frame unless `project_pose_to_goal_frame=False` is set in curobo.
             # For sticky grasping, we also need to ignore the object during motion planning because the fingers are already closed.
             yield from self._move_hand(
-                grasp_pose, motion_constraint=[1, 1, 1, 1, 1, 0], stop_on_contact=True, ignore_objects=[obj]
+                grasp_pose, motion_constraint=[1, 1, 1, 1, 1, 0], stop_on_ag=True, ignore_objects=[obj]
             )
         elif self.robot.grasping_mode == "assisted":
             indented_print("Assisted grasping: approach")
@@ -853,6 +856,7 @@ def _move_hand(
         self,
         target_pose,
         stop_on_contact=False,
+        stop_on_ag=False,
         motion_constraint=None,
         low_precision=False,
         lock_auxiliary_arm=False,
@@ -864,6 +868,7 @@ def _move_hand(
         Args:
             target_pose (Tuple[th.tensor, th.tensor]): target pose to reach for the default end-effector in the world frame
             stop_on_contact (bool): Whether to stop executing motion plan if contact is detected
+            stop_on_ag (bool): Whether to stop executing motion plan if assisted grasping is activated
             motion_constraint (MotionConstraint): Motion constraint for the motion
             low_precision (bool): Whether to use low precision for the motion
             lock_auxiliary_arm (bool): Whether to lock the other arm in place
@@ -906,7 +911,9 @@ def _move_hand(
         )
 
         indented_print(f"Plan has {len(q_traj)} steps")
-        yield from self._execute_motion_plan(q_traj, stop_on_contact=stop_on_contact, low_precision=low_precision)
+        yield from self._execute_motion_plan(
+            q_traj, stop_on_contact=stop_on_contact, stop_on_ag=stop_on_ag, low_precision=low_precision
+        )
 
     def _plan_joint_motion(
         self,
@@ -972,7 +979,13 @@ def _plan_joint_motion(
         return q_traj
 
     def _execute_motion_plan(
-        self, q_traj, stop_on_contact=False, ignore_failure=False, low_precision=False, ignore_physics=False
+        self,
+        q_traj,
+        stop_on_contact=False,
+        stop_on_ag=False,
+        ignore_failure=False,
+        low_precision=False,
+        ignore_physics=False,
     ):
         for i, joint_pos in enumerate(q_traj):
             # indented_print(f"Executing motion plan step {i + 1}/{len(q_traj)}")
@@ -992,6 +1005,12 @@ def _execute_motion_plan(
                 articulation_control_idx = th.cat(articulation_control_idx)
                 for j in range(m.MAX_STEPS_FOR_JOINT_MOTION):
                     # indented_print(f"Step {j + 1}/{m.MAX_STEPS_FOR_JOINT_MOTION}")
+
+                    # We need to call @q_to_action for every step because as the robot base moves, the same base joint_pos will be
+                    # converted to different actions, since HolonomicBaseJointController accepts an action in the robot local frame.
+                    action = self.robot.q_to_action(joint_pos)
+                    yield self._postprocess_action(action)
+
                     current_joint_pos = self.robot.get_joint_positions()
                     joint_pos_diff = joint_pos - current_joint_pos
                     base_joint_diff = joint_pos_diff[self.robot.base_control_idx]
@@ -1014,14 +1033,10 @@ def _execute_motion_plan(
 
                     collision_detected = detect_robot_collision_in_sim(self.robot)
                     if stop_on_contact and collision_detected:
-                        indented_print(f"Collision detected at step {i + 1}/{len(q_traj)}")
                         return
 
-                    # We need to call @q_to_action for every step because as the robot base moves, the same base joint_pos will be
-                    # converted to different actions, since HolonomicBaseJointController accepts an action in the robot local frame.
-                    action = self.robot.q_to_action(joint_pos)
-
-                    yield self._postprocess_action(action)
+                    if stop_on_ag and self._get_obj_in_hand() is not None:
+                        return
 
                 if not ignore_failure:
                     if not base_target_reached:
@@ -1334,13 +1349,13 @@ def _move_fingers_to_limit(self, limit_type):
         target_joint_positions = self._get_joint_position_with_fingers_at_limit(limit_type)
         action = self.robot.q_to_action(target_joint_positions)
         for _ in range(m.MAX_STEPS_FOR_GRASP_OR_RELEASE):
-            current_joint_positinos = self.robot.get_joint_positions()
-            if th.allclose(current_joint_positinos, target_joint_positions, atol=0.005):
+            yield self._postprocess_action(action)
+            current_joint_positions = self.robot.get_joint_positions()
+            if th.allclose(current_joint_positions, target_joint_positions, atol=m.JOINT_POS_DIFF_THRESHOLD):
                 break
             elif limit_type == "lower" and self._get_obj_in_hand() is not None:
                 # If we are grasping an object, we should stop when object is detected in hand
                 break
-            yield self._postprocess_action(action)
 
     def _execute_grasp(self):
         """

omnigibson/configs/fetch_behavior.yaml

@@ -40,6 +40,8 @@ scene:
 robots:
   - type: Fetch
     obs_modalities: [scan, rgb, depth]
+    include_sensor_names: null
+    exclude_sensor_names: null
     scale: 1.0
     self_collision: false
     action_normalize: true

omnigibson/configs/r1_primitives.yaml

@@ -32,6 +32,8 @@ scene:
 robots:
   - type: R1
     obs_modalities: [rgb]
+    include_sensor_names: null
+    exclude_sensor_names: null
     scale: 1.0
     self_collisions: true
     action_normalize: false

omnigibson/configs/robots/fetch.yaml

@@ -17,6 +17,8 @@ robot:
   self_collision: true
   grasping_mode: physical
   default_arm_pose: vertical
+  include_sensor_names: null
+  exclude_sensor_names: null
   sensor_config:
     VisionSensor:
       sensor_kwargs:

omnigibson/configs/robots/freight.yaml

@@ -10,6 +10,8 @@ robot:
   base_name: null
   scale: 1.0
   self_collision: false
+  include_sensor_names: null
+  exclude_sensor_names: null
   sensor_config:
     VisionSensor:
       sensor_kwargs:

omnigibson/configs/robots/husky.yaml

@@ -10,6 +10,8 @@ robot:
   base_name: null
   scale: 1.0
   self_collision: false
+  include_sensor_names: null
+  exclude_sensor_names: null
   sensor_config:
     VisionSensor:
       sensor_kwargs:

omnigibson/configs/robots/locobot.yaml

@@ -10,6 +10,8 @@ robot:
   base_name: null
   scale: 1.0
   self_collision: false
+  include_sensor_names: null
+  exclude_sensor_names: null
   sensor_config:
     VisionSensor:
       sensor_kwargs:

omnigibson/configs/robots/turtlebot.yaml

@@ -10,6 +10,8 @@ robot:
   base_name: null
   scale: 1.0
   self_collision: false
+  include_sensor_names: null
+  exclude_sensor_names: null
   sensor_config:
     VisionSensor:
       sensor_kwargs:

omnigibson/configs/tiago_primitives.yaml

@@ -33,6 +33,8 @@ scene:
 robots:
   - type: Tiago
     obs_modalities: [rgb]
+    include_sensor_names: null
+    exclude_sensor_names: null
     scale: 1.0
     self_collisions: true
     action_normalize: false

omnigibson/configs/turtlebot_nav.yaml

@@ -30,6 +30,8 @@ scene:
 robots:
   - type: Turtlebot
     obs_modalities: [scan, rgb, depth]
+    include_sensor_names: null
+    exclude_sensor_names: null
     scale: 1.0
     self_collision: false
     action_normalize: true

omnigibson/controllers/controller_base.py

@@ -4,9 +4,17 @@
 
 import torch as th
 
+from omnigibson.macros import create_module_macros
 from omnigibson.utils.backend_utils import _compute_backend as cb
 from omnigibson.utils.python_utils import Recreatable, Registerable, Serializable, assert_valid_key, classproperty
 
+# Create settings for this module
+m = create_module_macros(module_path=__file__)
+
+# Set default isaac kp / kd for controllers
+m.DEFAULT_ISAAC_KP = 1e7
+m.DEFAULT_ISAAC_KD = 1e5
+
 # Global dicts that will contain mappings
 REGISTERED_CONTROLLERS = dict()
 REGISTERED_LOCOMOTION_CONTROLLERS = dict()
@@ -76,6 +84,8 @@ def __init__(
         dof_idx,
         command_input_limits="default",
         command_output_limits="default",
+        isaac_kp=None,
+        isaac_kd=None,
     ):
         """
         Args:
@@ -99,6 +109,12 @@ def __init__(
                 then all inputted command values will be scaled from the input range to the output range.
                 If either is None, no scaling will be used. If "default", then this range will automatically be set
                 to the @control_limits entry corresponding to self.control_type
+            isaac_kp (None or float or Array[float]): If specified, stiffness gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers.
+                Should only be nonzero if self.control_type is position
+            isaac_kd (None or float or Array[float]): If specified, damping gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers
+                Should only be nonzero if self.control_type is position or velocity
         """
         # Store arguments
         self._control_freq = control_freq
@@ -154,6 +170,33 @@ def __init__(
             )
         )
 
+        # Set gains
+        if self.control_type == ControlType.POSITION:
+            # Set default kp / kd values if not specified
+            isaac_kp = m.DEFAULT_ISAAC_KP if isaac_kp is None else isaac_kp
+            isaac_kd = m.DEFAULT_ISAAC_KD if isaac_kd is None else isaac_kd
+        elif self.control_type == ControlType.VELOCITY:
+            # No kp should be specified, but kd should be
+            assert (
+                isaac_kp is None
+            ), f"Control type for controller {self.__class__.__name__} is VELOCITY, so no isaac_kp should be set!"
+            isaac_kd = m.DEFAULT_ISAAC_KP if isaac_kd is None else isaac_kd
+        elif self.control_type == ControlType.EFFORT:
+            # Neither kp nor kd should be specified
+            assert (
+                isaac_kp is None
+            ), f"Control type for controller {self.__class__.__name__} is EFFORT, so no isaac_kp should be set!"
+            assert (
+                isaac_kd is None
+            ), f"Control type for controller {self.__class__.__name__} is EFFORT, so no isaac_kd should be set!"
+        else:
+            raise ValueError(
+                f"Expected control type to be one of: [POSITION, VELOCITY, EFFORT], but got: {self.control_type}"
+            )
+
+        self._isaac_kp = None if isaac_kp is None else self.nums2array(isaac_kp, self.control_dim)
+        self._isaac_kd = None if isaac_kd is None else self.nums2array(isaac_kd, self.control_dim)
+
     def _generate_default_command_input_limits(self):
         """
         Generates default command input limits based on the control limits
@@ -510,6 +553,24 @@ def control_type(self):
         """
         raise NotImplementedError
 
+    @property
+    def isaac_kp(self):
+        """
+        Returns:
+            None or Array[float]: Stiffness gains that should be applied to the underlying Isaac joint motors.
+                None if not specified.
+        """
+        return self._isaac_kp
+
+    @property
+    def isaac_kd(self):
+        """
+        Returns:
+            None or Array[float]: Stiffness gains that should be applied to the underlying Isaac joint motors.
+                None if not specified.
+        """
+        return self._isaac_kd
+
     @property
     def command_input_limits(self):
         """

omnigibson/controllers/dd_controller.py

@@ -21,6 +21,8 @@ def __init__(
         dof_idx,
         command_input_limits="default",
         command_output_limits="default",
+        isaac_kp=None,
+        isaac_kd=None,
     ):
         """
         Args:
@@ -47,6 +49,12 @@ def __init__(
                 If either is None, no scaling will be used. If "default", then this range will automatically be set
                 to the maximum linear and angular velocities calculated from @wheel_radius, @wheel_axle_length, and
                 @control_limits velocity limits entry
+            isaac_kp (None or float or Array[float]): If specified, stiffness gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers.
+                Should only be nonzero if self.control_type is position
+            isaac_kd (None or float or Array[float]): If specified, damping gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers
+                Should only be nonzero if self.control_type is position or velocity
         """
         # Store internal variables
         self._wheel_radius = wheel_radius
@@ -76,6 +84,8 @@ def __init__(
             dof_idx=dof_idx,
             command_input_limits=command_input_limits,
             command_output_limits=command_output_limits,
+            isaac_kp=isaac_kp,
+            isaac_kd=isaac_kd,
         )
 
     def _update_goal(self, command, control_dict):

omnigibson/controllers/holonomic_base_joint_controller.py

@@ -25,6 +25,8 @@ def __init__(
         dof_idx,
         command_input_limits=None,
         command_output_limits=None,
+        isaac_kp=None,
+        isaac_kd=None,
         pos_kp=None,
         pos_damping_ratio=None,
         vel_kp=None,
@@ -55,6 +57,12 @@ def __init__(
                 then all inputted command values will be scaled from the input range to the output range.
                 If either is None, no scaling will be used. If "default", then this range will automatically be set
                 to the @control_limits entry corresponding to self.control_type
+            isaac_kp (None or float or Array[float]): If specified, stiffness gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers.
+                Should only be nonzero if self.control_type is position
+            isaac_kd (None or float or Array[float]): If specified, damping gains to apply to the underlying
+                isaac DOFs. Can either be a single number or a per-DOF set of numbers
+                Should only be nonzero if self.control_type is position or velocity
             pos_kp (None or float): If @motor_type is "position" and @use_impedances=True, this is the
                 proportional gain applied to the joint controller. If None, a default value will be used.
             pos_damping_ratio (None or float): If @motor_type is "position" and @use_impedances=True, this is the
@@ -78,6 +86,8 @@ def __init__(
             dof_idx=dof_idx,
             command_input_limits=command_input_limits,
             command_output_limits=command_output_limits,
+            isaac_kp=isaac_kp,
+            isaac_kd=isaac_kd,
             pos_kp=pos_kp,
             pos_damping_ratio=pos_damping_ratio,
             vel_kp=vel_kp,