Upgrade to isaac sim 4.2 & VR

omnigibson/controllers/ik_controller.py

@@ -1,5 +1,6 @@
 import math
 from collections.abc import Iterable
+from omnigibson.controllers.controller_base import _ControllerBackend, _ControllerTorchBackend, _ControllerNumpyBackend
 from omnigibson.controllers.controller_base import _controller_backend as cb
 
 from omnigibson.controllers import ControlType, ManipulationController
@@ -270,8 +271,8 @@ def _update_goal(self, command, control_dict):
             target_pos, target_quat = self.workspace_pose_limiter(target_pos, target_quat, control_dict)
 
         goal_dict = dict(
-            target_pos=target_pos,
-            target_quat=target_quat,
+            target_pos=cb.as_float32(target_pos),
+            target_ori_mat=cb.as_float32(cb.T.quat2mat(target_quat)),
         )
 
         return goal_dict
@@ -285,7 +286,7 @@ def compute_control(self, goal_dict, control_dict):
             goal_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped
                 goals necessary for controller computation. Must include the following keys:
                     target_pos: robot-frame (x,y,z) desired end effector position
-                    target_quat: robot-frame (x,y,z,w) desired end effector quaternion orientation
+                    target_ori_mat: robot-frame desired end effector quaternion orientation matrix
             control_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped
                 states necessary for controller computation. Must include the following keys:
                     joint_position: Array of current joint positions
@@ -299,37 +300,23 @@ def compute_control(self, goal_dict, control_dict):
         Returns:
             Array[float]: outputted (non-clipped!) velocity control signal to deploy
         """
-        # Grab important info from control dict
-        pos_relative = control_dict[f"{self.task_name}_pos_relative"]
-        quat_relative = control_dict[f"{self.task_name}_quat_relative"]
-        target_pos = goal_dict["target_pos"]
-        target_quat = goal_dict["target_quat"]
-
         # Calculate and return IK-backed out joint angles
-        current_joint_pos = control_dict["joint_position"][self.dof_idx]
-
-        # If the delta is really small, we just keep the current joint position. This avoids joint
-        # drift caused by IK solver inaccuracy even when zero delta actions are provided.
-        if cb.allclose(pos_relative, target_pos, atol=1e-4) and cb.allclose(quat_relative, target_quat, atol=1e-4):
-            target_joint_pos = current_joint_pos
-        else:
-            # Compute the pose error. Note that this is computed NOT in the EEF frame but still
-            # in the base frame.
-            pos_err = target_pos - pos_relative
-            ori_err = cb.T.orientation_error(cb.T.quat2mat(target_quat), cb.T.quat2mat(quat_relative))
-            err = cb.cat([pos_err, ori_err])
-
-            # Use the jacobian to compute a local approximation
-            j_eef = control_dict[f"{self.task_name}_jacobian_relative"][:, self.dof_idx]
-            j_eef_pinv = cb.pinv(j_eef)
-            delta_j = j_eef_pinv @ err
-            target_joint_pos = current_joint_pos + delta_j
-
-            # Clip values to be within the joint limits
-            target_joint_pos = target_joint_pos.clip(
-                min=self._control_limits[ControlType.get_type("position")][0][self.dof_idx],
-                max=self._control_limits[ControlType.get_type("position")][1][self.dof_idx],
-            )
+        q = control_dict["joint_position"][self.dof_idx]
+        j_eef = control_dict[f"{self.task_name}_jacobian_relative"][:, self.dof_idx]
+        ee_pos = control_dict[f"{self.task_name}_pos_relative"]
+        ee_quat = control_dict[f"{self.task_name}_quat_relative"]
+
+        # Calculate desired joint positions
+        target_joint_pos = cb.compute_ik_qpos(
+            q=q,
+            j_eef=j_eef,
+            ee_pos=cb.as_float32(ee_pos),
+            ee_mat=cb.as_float32(cb.T.quat2mat(ee_quat)),
+            goal_pos=goal_dict["target_pos"],
+            goal_ori_mat=goal_dict["target_ori_mat"],
+            q_lower_limit=self._control_limits[ControlType.get_type("position")][0][self.dof_idx],
+            q_upper_limit=self._control_limits[ControlType.get_type("position")][1][self.dof_idx],
+        )
 
         # Optionally pass through smoothing filter for better stability
         if self.control_filter is not None:
@@ -340,9 +327,13 @@ def compute_control(self, goal_dict, control_dict):
 
     def compute_no_op_goal(self, control_dict):
         # No-op is maintaining current pose
+        target_pos = cb.copy(control_dict[f"{self.task_name}_pos_relative"])
+        target_quat = cb.copy(control_dict[f"{self.task_name}_quat_relative"])
+
+        # Convert quat into eef ori mat
         return dict(
-            target_pos=control_dict[f"{self.task_name}_pos_relative"],
-            target_quat=control_dict[f"{self.task_name}_quat_relative"],
+            target_pos=cb.as_float32(target_pos),
+            target_ori_mat=cb.as_float32(cb.T.quat2mat(target_quat)),
         )
 
     def _compute_no_op_action(self, control_dict):
@@ -370,9 +361,77 @@ def _compute_no_op_action(self, control_dict):
     def _get_goal_shapes(self):
         return dict(
             target_pos=(3,),
-            target_quat=(4,),
+            target_ori_mat=(3, 3),
         )
 
     @property
     def command_dim(self):
         return IK_MODE_COMMAND_DIMS[self.mode]
+
+
+import torch as th
+import omnigibson.utils.transform_utils as TT
+@th.jit.script
+def _compute_ik_qpos_torch(
+    q: th.Tensor,
+    j_eef: th.Tensor,
+    ee_pos: th.Tensor,
+    ee_mat: th.Tensor,
+    goal_pos: th.Tensor,
+    goal_ori_mat: th.Tensor,
+    q_lower_limit: th.Tensor,
+    q_upper_limit: th.Tensor,
+):
+    # Compute the pose error. Note that this is computed NOT in the EEF frame but still
+    # in the base frame.
+    pos_err = goal_pos - ee_pos
+    ori_err = TT.orientation_error(goal_ori_mat, ee_mat)
+    err = th.cat([pos_err, ori_err])
+
+    # Use the jacobian to compute a local approximation
+    j_eef_pinv = th.linalg.pinv(j_eef)
+    delta_j = j_eef_pinv @ err
+    target_joint_pos = q + delta_j
+
+    # Clip values to be within the joint limits
+    return target_joint_pos.clip(
+        min=q_lower_limit,
+        max=q_upper_limit,
+    )
+
+
+import numpy as np
+from numba import jit
+import omnigibson.utils.transform_utils_np as NT
+# Use numba since faster
+@jit(nopython=True)
+def _compute_ik_qpos_numpy(
+    q,
+    j_eef,
+    ee_pos,
+    ee_mat,
+    goal_pos,
+    goal_ori_mat,
+    q_lower_limit,
+    q_upper_limit,
+):
+    # Compute the pose error. Note that this is computed NOT in the EEF frame but still
+    # in the base frame.
+    pos_err = goal_pos - ee_pos
+    ori_err = NT.orientation_error(goal_ori_mat, ee_mat).astype(np.float32)
+    err = np.concatenate((pos_err, ori_err))
+
+    # Use the jacobian to compute a local approximation
+    j_eef_pinv = np.linalg.pinv(j_eef)
+    delta_j = j_eef_pinv @ err
+    target_joint_pos = q + delta_j
+
+    # Clip values to be within the joint limits
+    return target_joint_pos.clip(q_lower_limit, q_upper_limit)
+
+
+# Set these as part of the backend values
+setattr(_ControllerBackend, "compute_ik_qpos", None)
+setattr(_ControllerTorchBackend, "compute_ik_qpos", _compute_ik_qpos_torch)
+setattr(_ControllerNumpyBackend, "compute_ik_qpos", _compute_ik_qpos_numpy)
+

omnigibson/controllers/joint_controller.py

@@ -1,5 +1,6 @@
 import math
 
+from omnigibson.controllers.controller_base import _ControllerBackend, _ControllerTorchBackend, _ControllerNumpyBackend
 from omnigibson.controllers.controller_base import _controller_backend as cb
 from omnigibson.controllers import (
     ControlType,
@@ -210,9 +211,7 @@ def compute_control(self, goal_dict, control_dict):
             else:  # effort
                 u = target
 
-            dof_idxs_mat = cb.meshgrid(self.dof_idx, self.dof_idx)
-            mm = control_dict["mass_matrix"][dof_idxs_mat]
-            u = mm @ u
+            u = cb.compute_joint_torques(u, control_dict["mass_matrix"], self.dof_idx)
 
             # Add gravity compensation
             if self._use_gravity_compensation:
@@ -284,3 +283,57 @@ def control_type(self):
     @property
     def command_dim(self):
         return len(self.dof_idx)
+
+
+import torch as th
+@th.compile
+def _compute_joint_torques_torch(
+    u: th.Tensor,
+    mm: th.Tensor,
+    dof_idx: th.Tensor,
+):
+    dof_idxs_mat = th.meshgrid(dof_idx, dof_idx, indexing="xy")
+    return mm[dof_idxs_mat] @ u
+
+
+
+import numpy as np
+from numba import jit
+# Use numba since faster
+@jit(nopython=True)
+def numba_ix(arr, rows, cols):
+    """
+    Numba compatible implementation of arr[np.ix_(rows, cols)] for 2D arrays.
+
+    Implementation from:
+    https://github.com/numba/numba/issues/5894#issuecomment-974701551
+
+    :param arr: 2D array to be indexed
+    :param rows: Row indices
+    :param cols: Column indices
+    :return: 2D array with the given rows and columns of the input array
+    """
+    one_d_index = np.zeros(len(rows) * len(cols), dtype=np.int32)
+    for i, r in enumerate(rows):
+        start = i * len(cols)
+        one_d_index[start: start + len(cols)] = cols + arr.shape[1] * r
+
+    arr_1d = arr.reshape((arr.shape[0] * arr.shape[1], 1))
+    slice_1d = np.take(arr_1d, one_d_index)
+    return slice_1d.reshape((len(rows), len(cols)))
+
+
+@jit(nopython=True)
+def _compute_joint_torques_numpy(
+    u,
+    mm,
+    dof_idx,
+):
+    return numba_ix(mm, dof_idx, dof_idx) @ u
+
+
+# Set these as part of the backend values
+setattr(_ControllerBackend, "compute_joint_torques", None)
+setattr(_ControllerTorchBackend, "compute_joint_torques", _compute_joint_torques_torch)
+setattr(_ControllerNumpyBackend, "compute_joint_torques", _compute_joint_torques_numpy)
+

omnigibson/controllers/osc_controller.py

@@ -362,7 +362,7 @@ def compute_control(self, goal_dict, control_dict):
             goal_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped
                 goals necessary for controller computation. Must include the following keys:
                     target_pos: robot-frame (x,y,z) desired end effector position
-                    target_quat: robot-frame (x,y,z,w) desired end effector quaternion orientation
+                    target_ori_mat: robot-frame desired end effector quaternion orientation matrix
             control_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped
                 states necessary for controller computation. Must include the following keys:
                     joint_position: Array of current joint positions

omnigibson/objects/object_base.py

@@ -237,7 +237,7 @@ def _initialize(self):
                 "emissive_intensity": material.emissive_intensity,
             }
 
-    @property
+    @cached_property
     def articulation_root_path(self):
         has_articulated_joints, has_fixed_joints = self.n_joints > 0, self.n_fixed_joints > 0
         if self.kinematic_only or ((not has_articulated_joints) and (not has_fixed_joints)):

omnigibson/prims/cloth_prim.py

@@ -9,6 +9,7 @@
 
 import math
 from collections.abc import Iterable
+from functools import cached_property
 
 import torch as th
 
@@ -160,7 +161,7 @@ def n_particles(self):
         """
         return self._n_particles
 
-    @property
+    @cached_property
     def kinematic_only(self):
         """
         Returns:

omnigibson/prims/entity_prim.py

@@ -1399,7 +1399,7 @@ def self_collisions(self, flag):
             self.articulation_root_path, "physxArticulation:enabledSelfCollisions", flag
         )
 
-    @property
+    @cached_property
     def kinematic_only(self):
         """
         Returns:

omnigibson/prims/joint_prim.py

@@ -75,6 +75,8 @@ def __init__(
         self._joint_type = None
         self._control_type = None
         self._driven = None
+        self._body0 = None
+        self._body1 = None
 
         # The following values will only be valid if this joint is part of an articulation
         self._n_dof = None  # The number of degrees of freedom this joint provides
@@ -232,8 +234,10 @@ def body0(self):
             None or str: Absolute prim path to the body prim to set as this joint's parent link, or None if there is
                 no body0 specified.
         """
-        targets = self._prim.GetRelationship("physics:body0").GetTargets()
-        return targets[0].__str__() if len(targets) > 0 else None
+        if self._body0 is None:
+            targets = self._prim.GetRelationship("physics:body0").GetTargets()
+            self._body0 = targets[0].__str__()
+        return self._body0
 
     @body0.setter
     def body0(self, body0):
@@ -246,6 +250,7 @@ def body0(self, body0):
         # Make sure prim path is valid
         assert lazy.omni.isaac.core.utils.prims.is_prim_path_valid(body0), f"Invalid body0 path specified: {body0}"
         self._prim.GetRelationship("physics:body0").SetTargets([lazy.pxr.Sdf.Path(body0)])
+        self._body0 = None
 
     @property
     def body1(self):
@@ -256,8 +261,10 @@ def body1(self):
             None or str: Absolute prim path to the body prim to set as this joint's child link, or None if there is
                 no body1 specified.
         """
-        targets = self._prim.GetRelationship("physics:body1").GetTargets()
-        return targets[0].__str__()
+        if self._body1 is None:
+            targets = self._prim.GetRelationship("physics:body1").GetTargets()
+            self._body1 = targets[0].__str__()
+        return self._body1
 
     @body1.setter
     def body1(self, body1):
@@ -270,6 +277,7 @@ def body1(self, body1):
         # Make sure prim path is valid
         assert lazy.omni.isaac.core.utils.prims.is_prim_path_valid(body1), f"Invalid body1 path specified: {body1}"
         self._prim.GetRelationship("physics:body1").SetTargets([lazy.pxr.Sdf.Path(body1)])
+        self._body1 = None
 
     @property
     def local_orientation(self):

omnigibson/robots/a1.py

@@ -1,4 +1,5 @@
 import os
+from functools import cached_property
 
 import torch as th
 
@@ -167,23 +168,23 @@ def _default_joint_pos(self):
     def finger_lengths(self):
         return {self.default_arm: 0.087}
 
-    @property
+    @cached_property
     def arm_link_names(self):
         return {self.default_arm: [f"arm_seg{i+1}" for i in range(5)]}
 
-    @property
+    @cached_property
     def arm_joint_names(self):
         return {self.default_arm: [f"arm_joint{i+1}" for i in range(6)]}
 
-    @property
+    @cached_property
     def eef_link_names(self):
         return {self.default_arm: self._eef_link_names}
 
-    @property
+    @cached_property
     def finger_link_names(self):
         return {self.default_arm: self._finger_link_names}
 
-    @property
+    @cached_property
     def finger_joint_names(self):
         return {self.default_arm: self._finger_joint_names}
 

omnigibson/robots/active_camera_robot.py

@@ -115,7 +115,7 @@ def _default_controller_config(self):
 
         return cfg
 
-    @property
+    @cached_property
     @abstractmethod
     def camera_joint_names(self):
         """