Merge pull request #797 from StanfordVL/feat/auto-robot-idx

Feat/auto robot idx

omnigibson/envs/env_base.py

@@ -594,14 +594,9 @@ def step(self, action):
                 - bool: truncated, i.e. whether this episode ended due to a time limit etc.
                 - dict: info, i.e. dictionary with any useful information
         """
-        try:
-            self._pre_step(action)
-            og.sim.step()
-            return self._post_step(action)
-        except:
-            raise ValueError(
-                f"Failed to execute environment step {self._current_step} in episode {self._current_episode}"
-            )
+        self._pre_step(action)
+        og.sim.step()
+        return self._post_step(action)
 
     def render(self):
         """Render the environment for debug viewing."""

omnigibson/objects/controllable_object.py

@@ -148,14 +148,6 @@ def _initialize(self):
         self.reset()
         self.keep_still()
 
-        # If we haven't already created a physics callback, do so now so control gets updated every sim step
-        callback_name = f"{self.name}_controller_callback"
-        if not og.sim.physics_callback_exists(callback_name=callback_name):
-            og.sim.add_physics_callback(
-                callback_name=callback_name,
-                callback_fn=lambda x: self.step(),
-            )
-
     def load(self, scene):
         # Run super first
         prim = super().load(scene)
@@ -559,7 +551,7 @@ def get_control_dict(self):
         # TODO: Move gravity force computation dummy to this class instead of BaseRobot
         fcns["gravity_force"] = lambda: (
             ControllableObjectViewAPI.get_generalized_gravity_forces(self.articulation_root_path)
-            if not self.fixed_base
+            if self.fixed_base
             else DummyControllableObjectViewAPI.get_generalized_gravity_forces(self._dummy.articulation_root_path)
         )
         fcns["cc_force"] = lambda: ControllableObjectViewAPI.get_coriolis_and_centrifugal_forces(

omnigibson/robots/__init__.py

@@ -9,6 +9,7 @@
 from omnigibson.robots.locomotion_robot import LocomotionRobot
 from omnigibson.robots.manipulation_robot import ManipulationRobot
 from omnigibson.robots.robot_base import REGISTERED_ROBOTS, BaseRobot
+from omnigibson.robots.stretch import Stretch
 from omnigibson.robots.tiago import Tiago
 from omnigibson.robots.turtlebot import Turtlebot
 from omnigibson.robots.two_wheel_robot import TwoWheelRobot

omnigibson/robots/active_camera_robot.py

@@ -116,12 +116,23 @@ def _default_controller_config(self):
 
     @property
     @abstractmethod
+    def camera_joint_names(self):
+        """
+        Returns:
+            list: Array of joint names corresponding to this robot's camera joints.
+
+                Note: the ordering within the list is assumed to be intentional, and is
+                directly used to define the set of corresponding control idxs.
+        """
+        raise NotImplementedError
+
+    @property
     def camera_control_idx(self):
         """
         Returns:
             n-array: Indices in low-level control vector corresponding to camera joints.
         """
-        raise NotImplementedError
+        return np.array([list(self.joints.keys()).index(name) for name in self.camera_joint_names])
 
     @classproperty
     def _do_not_register_classes(cls):

omnigibson/robots/behavior_robot.py

@@ -142,10 +142,6 @@ def __init__(
     def usd_path(self):
         return os.path.join(gm.ASSET_PATH, "models/behavior_robot/usd/BehaviorRobot.usd")
 
-    @property
-    def model_name(self):
-        return "BehaviorRobot"
-
     @classproperty
     def n_arms(cls):
         return 2
@@ -178,6 +174,10 @@ def finger_link_names(self):
     def base_joint_names(self):
         return [f"base_{component}_joint" for component in m.COMPONENT_SUFFIXES]
 
+    @property
+    def camera_joint_names(self):
+        return [f"head_{component}_joint" for component in m.COMPONENT_SUFFIXES]
+
     @property
     def arm_joint_names(self):
         """The head counts as a arm since it has the same 33 joint configuration"""
@@ -201,29 +201,6 @@ def finger_joint_names(self):
             for arm in self.arm_names
         }
 
-    @property
-    def base_control_idx(self):
-        joints = list(self.joints.keys())
-        return [joints.index(joint) for joint in self.base_joint_names]
-
-    @property
-    def arm_control_idx(self):
-        joints = list(self.joints.keys())
-        return {
-            arm: [joints.index(f"{arm}_{component}_joint") for component in m.COMPONENT_SUFFIXES]
-            for arm in self.arm_names
-        }
-
-    @property
-    def gripper_control_idx(self):
-        joints = list(self.joints.values())
-        return {arm: [joints.index(joint) for joint in arm_joints] for arm, arm_joints in self.finger_joints.items()}
-
-    @property
-    def camera_control_idx(self):
-        joints = list(self.joints.keys())
-        return [joints.index(f"head_{component}_joint") for component in m.COMPONENT_SUFFIXES]
-
     @property
     def _default_joint_pos(self):
         return np.zeros(self.n_joints)

omnigibson/robots/fetch.py

@@ -158,10 +158,6 @@ def __init__(
             **kwargs,
         )
 
-    @property
-    def model_name(self):
-        return "Fetch"
-
     @property
     def tucked_default_joint_pos(self):
         return np.array(
@@ -364,37 +360,13 @@ def assisted_grasp_end_points(self):
             ]
         }
 
-    @property
-    def base_control_idx(self):
-        """
-        Returns:
-            n-array: Indices in low-level control vector corresponding to [Left, Right] wheel joints.
-        """
-        return np.array([0, 1])
-
     @property
     def trunk_control_idx(self):
         """
         Returns:
-            n-array: Indices in low-level control vector corresponding to trunk joint.
-        """
-        return np.array([2])
-
-    @property
-    def camera_control_idx(self):
-        """
-        Returns:
-            n-array: Indices in low-level control vector corresponding to [tilt, pan] camera joints.
+            n-array: Indices in low-level control vector corresponding to trunk joints.
         """
-        return np.array([3, 5])
-
-    @property
-    def arm_control_idx(self):
-        return {self.default_arm: np.array([4, 6, 7, 8, 9, 10, 11])}
-
-    @property
-    def gripper_control_idx(self):
-        return {self.default_arm: np.array([12, 13])}
+        return np.array([list(self.joints.keys()).index(name) for name in self.trunk_joint_names])
 
     @property
     def disabled_collision_pairs(self):
@@ -419,6 +391,18 @@ def disabled_collision_pairs(self):
             ["wrist_roll_link", "gripper_link"],
         ]
 
+    @property
+    def base_joint_names(self):
+        return ["l_wheel_joint", "r_wheel_joint"]
+
+    @property
+    def camera_joint_names(self):
+        return ["head_pan_joint", "head_tilt_joint"]
+
+    @property
+    def trunk_joint_names(self):
+        return ["torso_lift_joint"]
+
     @property
     def manipulation_link_names(self):
         return [
@@ -455,7 +439,6 @@ def arm_link_names(self):
     def arm_joint_names(self):
         return {
             self.default_arm: [
-                "torso_lift_joint",
                 "shoulder_pan_joint",
                 "shoulder_lift_joint",
                 "upperarm_roll_joint",

omnigibson/robots/franka.py

@@ -207,6 +207,7 @@ def __init__(
 
     @property
     def model_name(self):
+        # Override based on specified Franka variant
         return self._model_name
 
     @property
@@ -250,18 +251,6 @@ def _default_joint_pos(self):
     def finger_lengths(self):
         return {self.default_arm: 0.1}
 
-    @property
-    def arm_control_idx(self):
-        return {self.default_arm: np.arange(7)}
-
-    @property
-    def gripper_control_idx(self):
-        return {
-            self.default_arm: np.array(
-                [list(self.joints.keys()).index(name) for name in self.finger_joint_names[self.default_arm]]
-            )
-        }
-
     @property
     def arm_link_names(self):
         return {self.default_arm: [f"panda_link{i}" for i in range(8)]}

omnigibson/robots/franka_mounted.py

@@ -13,10 +13,6 @@ class FrankaMounted(FrankaPanda):
     The Franka Emika Panda robot mounted on a custom chassis with a custom gripper
     """
 
-    @property
-    def model_name(self):
-        return "FrankaMounted"
-
     @property
     def controller_order(self):
         return ["arm_{}".format(self.default_arm), "gripper_{}".format(self.default_arm)]

omnigibson/robots/freight.py

@@ -13,10 +13,6 @@ class Freight(TwoWheelRobot):
     Uses joint velocity control
     """
 
-    @property
-    def model_name(self):
-        return "Freight"
-
     @property
     def wheel_radius(self):
         return 0.0613
@@ -26,12 +22,8 @@ def wheel_axle_length(self):
         return 0.372
 
     @property
-    def base_control_idx(self):
-        """
-        Returns:
-            n-array: Indices in low-level control vector corresponding to [Left, Right] wheel joints.
-        """
-        return np.array([0, 1])
+    def base_joint_names(self):
+        return ["r_wheel_joint", "l_wheel_joint"]
 
     @property
     def _default_joint_pos(self):

omnigibson/robots/husky.py

@@ -24,8 +24,8 @@ def wheel_axle_length(self):
         return 0.670
 
     @property
-    def base_control_idx(self):
-        return np.array([0, 1, 2, 3])
+    def base_joint_names(self):
+        return ["front_left_wheel", "front_right_wheel", "rear_left_wheel", "rear_right_wheel"]
 
     @property
     def _default_joint_pos(self):

omnigibson/robots/locobot.py

@@ -12,10 +12,6 @@ class Locobot(TwoWheelRobot):
     Reference: https://www.trossenrobotics.com/locobot-pyrobot-ros-rover.aspx
     """
 
-    @property
-    def model_name(self):
-        return "Locobot"
-
     @property
     def wheel_radius(self):
         return 0.038
@@ -25,12 +21,8 @@ def wheel_axle_length(self):
         return 0.230
 
     @property
-    def base_control_idx(self):
-        """
-        Returns:
-            n-array: Indices in low-level control vector corresponding to [Left, Right] wheel joints.
-        """
-        return np.array([1, 0])
+    def base_joint_names(self):
+        return ["wheel_right_joint", "wheel_left_joint"]
 
     @property
     def _default_joint_pos(self):

omnigibson/robots/locomotion_robot.py

@@ -193,12 +193,23 @@ def base_action_idx(self):
 
     @property
     @abstractmethod
+    def base_joint_names(self):
+        """
+        Returns:
+            list: Array of joint names corresponding to this robot's base joints (e.g.: wheels).
+
+                Note: the ordering within the list is assumed to be intentional, and is
+                directly used to define the set of corresponding control idxs.
+        """
+        raise NotImplementedError
+
+    @property
     def base_control_idx(self):
         """
         Returns:
             n-array: Indices in low-level control vector corresponding to base joints.
         """
-        raise NotImplementedError
+        return np.array([list(self.joints.keys()).index(name) for name in self.base_joint_names])
 
     @classproperty
     def _do_not_register_classes(cls):