Merge pull request #608 from StanfordVL/feat/controller-optimization

Optimize Controller Logic

omnigibson/objects/controllable_object.py

@@ -512,9 +512,9 @@ def get_control_dict(self):
         fcns["joint_position"] = lambda: self.get_joint_positions(normalized=False)
         fcns["joint_velocity"] = lambda: self.get_joint_velocities(normalized=False)
         fcns["joint_effort"] = lambda: self.get_joint_efforts(normalized=False)
-        fcns["mass_matrix"] = self.get_mass_matrix
-        fcns["gravity_force"] = self.get_generalized_gravity_forces
-        fcns["cc_force"] = self.get_coriolis_and_centrifugal_forces
+        fcns["mass_matrix"] = lambda: self.get_mass_matrix(clone=False)
+        fcns["gravity_force"] = lambda: self.get_generalized_gravity_forces(clone=False)
+        fcns["cc_force"] = lambda: self.get_coriolis_and_centrifugal_forces(clone=False)
 
         return fcns
 

omnigibson/objects/object_base.py

@@ -165,6 +165,12 @@ def _post_load(self):
                     body1=f"{self._prim_path}/{self._root_link_name}",
                 )
 
+            # Delete n_fixed_joints cached property if it exists since the number of fixed joints has now changed
+            # See https://stackoverflow.com/questions/59899732/python-cached-property-how-to-delete and
+            # https://docs.python.org/3/library/functools.html#functools.cached_property
+            if "n_fixed_joints" in self.__dict__:
+                del self.n_fixed_joints
+
         # Set visibility
         if "visible" in self._load_config and self._load_config["visible"] is not None:
             self.visible = self._load_config["visible"]

omnigibson/prims/entity_prim.py

@@ -1,5 +1,6 @@
 import numpy as np
 import networkx as nx
+from functools import cached_property
 
 import omnigibson as og
 import omnigibson.lazy as lazy
@@ -420,7 +421,7 @@ def n_joints(self):
             int: Number of joints owned by this articulation
         """
         if self.initialized:
-            num = len(list(self._joints.keys()))
+            num = len(self._joints)
         else:
             # Manually iterate over all links and check for any joints that are not fixed joints!
             num = 0
@@ -433,7 +434,7 @@ def n_joints(self):
                     num += 1
         return num
 
-    @property
+    @cached_property
     def n_fixed_joints(self):
         """
         Returns:
@@ -616,13 +617,12 @@ def set_joint_positions(self, positions, indices=None, normalized=False, drive=F
             positions = self._denormalize_positions(positions=positions, indices=indices)
 
         # Set the DOF states
-        if not drive:
+        if drive:
+            self._articulation_view.set_joint_position_targets(positions, joint_indices=indices)
+        else:
             self._articulation_view.set_joint_positions(positions, joint_indices=indices)
             BoundingBoxAPI.clear()
 
-        # Also set the target
-        self._articulation_view.set_joint_position_targets(positions, joint_indices=indices)
-
     def set_joint_velocities(self, velocities, indices=None, normalized=False, drive=False):
         """
         Set the joint velocities (both actual value and target values) in simulation. Note: only works if the simulator
@@ -648,12 +648,11 @@ def set_joint_velocities(self, velocities, indices=None, normalized=False, drive
             velocities = self._denormalize_velocities(velocities=velocities, indices=indices)
 
         # Set the DOF states
-        if not drive:
+        if drive:
+            self._articulation_view.set_joint_velocity_targets(velocities, joint_indices=indices)
+        else:
             self._articulation_view.set_joint_velocities(velocities, joint_indices=indices)
 
-        # Also set the target
-        self._articulation_view.set_joint_velocity_targets(velocities, joint_indices=indices)
-
     def set_joint_efforts(self, efforts, indices=None, normalized=False):
         """
         Set the joint efforts (both actual value and target values) in simulation. Note: only works if the simulator
@@ -1247,49 +1246,64 @@ def aabb(self):
 
         return aabb_lo, aabb_hi
 
-    def get_coriolis_and_centrifugal_forces(self):
+    def get_coriolis_and_centrifugal_forces(self, clone=True):
         """
+        Args:
+            clone (bool): Whether to clone the underlying tensor buffer or not
+
         Returns:
             n-array: (N,) shaped per-DOF coriolis and centrifugal forces experienced by the entity, if articulated
         """
         assert self.articulated, "Cannot get coriolis and centrifugal forces for non-articulated entity!"
-        return self._articulation_view.get_coriolis_and_centrifugal_forces().reshape(self.n_dof)
+        return self._articulation_view.get_coriolis_and_centrifugal_forces(clone=clone).reshape(self.n_dof)
 
-    def get_generalized_gravity_forces(self):
+    def get_generalized_gravity_forces(self, clone=True):
         """
+        Args:
+            clone (bool): Whether to clone the underlying tensor buffer or not
+
         Returns:
             n-array: (N, N) shaped per-DOF gravity forces, if articulated
         """
         assert self.articulated, "Cannot get generalized gravity forces for non-articulated entity!"
-        return self._articulation_view.get_generalized_gravity_forces().reshape(self.n_dof)
+        return self._articulation_view.get_generalized_gravity_forces(clone=clone).reshape(self.n_dof)
 
-    def get_mass_matrix(self):
+    def get_mass_matrix(self, clone=True):
         """
+        Args:
+            clone (bool): Whether to clone the underlying tensor buffer or not
+
         Returns:
             n-array: (N, N) shaped per-DOF mass matrix, if articulated
         """
         assert self.articulated, "Cannot get mass matrix for non-articulated entity!"
-        return self._articulation_view.get_mass_matrices().reshape(self.n_dof, self.n_dof)
+        return self._articulation_view.get_mass_matrices(clone=clone).reshape(self.n_dof, self.n_dof)
 
-    def get_jacobian(self):
+    def get_jacobian(self, clone=True):
         """
+        Args:
+            clone (bool): Whether to clone the underlying tensor buffer or not
+
         Returns:
             n-array: (N_links - 1 [+ 1], 6, N_dof [+ 6]) shaped per-link jacobian, if articulated. Note that the first
                 dimension is +1 and the final dimension is +6 if the entity does not have a fixed base
                 (i.e.: there is an additional "floating" joint tying the robot to the world frame)
         """
         assert self.articulated, "Cannot get jacobian for non-articulated entity!"
-        return self._articulation_view.get_jacobians().squeeze(axis=0)
+        return self._articulation_view.get_jacobians(clone=clone).squeeze(axis=0)
 
-    def get_relative_jacobian(self):
+    def get_relative_jacobian(self, clone=True):
         """
+        Args:
+            clone (bool): Whether to clone the underlying tensor buffer or not
+
         Returns:
             n-array: (N_links - 1 [+ 1], 6, N_dof [+ 6]) shaped per-link relative jacobian, if articulated (expressed in
                 this entity's base frame). Note that the first dimension is +1 and the final dimension is +6 if the
                 entity does not have a fixed base (i.e.: there is an additional "floating" joint tying the robot to
                 the world frame)
         """
-        jac = self.get_jacobian()
+        jac = self.get_jacobian(clone=clone)
         ori_t = T.quat2mat(self.get_orientation()).T.astype(np.float32)
         tf = np.zeros((1, 6, 6), dtype=np.float32)
         tf[:, :3, :3] = ori_t

omnigibson/robots/manipulation_robot.py

@@ -404,7 +404,7 @@ def _add_arm_control_dict(self, fcns, arm):
         # not have a fixed base (i.e.: the 6DOF --> "floating" joint)
         # see self.get_relative_jacobian() for more info
         eef_link_idx = self._articulation_view.get_body_index(self.eef_links[arm].body_name)
-        fcns[f"eef_{arm}_jacobian_relative"] = lambda: self.get_relative_jacobian()[eef_link_idx, :, -self.n_joints:]
+        fcns[f"eef_{arm}_jacobian_relative"] = lambda: self.get_relative_jacobian(clone=False)[eef_link_idx, :, -self.n_joints:]
 
     def _get_proprioception_dict(self):
         dic = super()._get_proprioception_dict()

omnigibson/robots/robot_base.py

@@ -271,6 +271,22 @@ def _validate_configuration(self):
         """
         pass
 
+    def step(self):
+        # Skip this step if our articulation view is not valid
+        if self._articulation_view_direct is None or not self._articulation_view_direct.initialized:
+            return
+
+        # Before calling super, update the dummy robot's kinematic state based on this robot's kinematic state
+        # This is done prior to any state getter calls, since setting kinematic state results in physx backend
+        # having to re-fetch tensorized state.
+        # We do this so we have more optimal runtime performance
+        if self._use_dummy:
+            self._dummy.set_joint_positions(self.get_joint_positions())
+            self._dummy.set_joint_velocities(self.get_joint_velocities())
+            self._dummy.set_position_orientation(*self.get_position_orientation())
+
+        super().step()
+
     def get_obs(self):
         """
         Grabs all observations from the robot. This is keyword-mapped based on each observation modality
@@ -471,16 +487,10 @@ def teleop_data_to_action(self, teleop_data):
         """
         return np.zeros(self.action_dim)
 
-    def get_generalized_gravity_forces(self):
+    def get_generalized_gravity_forces(self, clone=True):
         # Override method based on whether we're using a dummy or not
-        if self._use_dummy:
-            # Update dummy pose and calculate values
-            self._dummy.set_joint_positions(self.get_joint_positions())
-            self._dummy.set_joint_velocities(self.get_joint_velocities())
-            self._dummy.set_position_orientation(*self.get_position_orientation())
-            return self._dummy.get_generalized_gravity_forces()
-        else:
-            return super().get_generalized_gravity_forces()
+        return self._dummy.get_generalized_gravity_forces(clone=clone) \
+            if self._use_dummy else super().get_generalized_gravity_forces(clone=clone)
 
     @property
     def sensors(self):