Torch mini followup - bug fixes

.github/workflows/tests.yml

@@ -20,6 +20,8 @@ jobs:
       fail-fast: false
       matrix:
         test_file:
+          - test_controllers
+          - test_data_collection
           - test_dump_load_states
           - test_envs
           - test_multiple_envs

omnigibson/action_primitives/starter_semantic_action_primitives.py

@@ -1112,7 +1112,7 @@ def _move_hand_direct_ik(
         target_pose,
         stop_on_contact=False,
         ignore_failure=False,
-        pos_thresh=0.04,
+        pos_thresh=0.02,
         ori_thresh=0.4,
         in_world_frame=True,
         stop_if_stuck=False,
@@ -1430,16 +1430,8 @@ def _empty_action(self):
         action = th.zeros(self.robot.action_dim)
         for name, controller in self.robot._controllers.items():
             action_idx = self.robot.controller_action_idx[name]
-            no_op_goal = controller.compute_no_op_goal(self.robot.get_control_dict())
-
-            if self.robot._controller_config[name]["name"] == "InverseKinematicsController":
-                assert (
-                    self.robot._controller_config["arm_" + self.arm]["mode"] == "pose_absolute_ori"
-                ), "Controller must be in pose_absolute_ori mode"
-                # convert quaternion to axis-angle representation for control input
-                no_op_goal["target_quat"] = T.quat2axisangle(no_op_goal["target_quat"])
-
-            action[action_idx] = th.cat(list(no_op_goal.values()))
+            no_op_action = controller.compute_no_op_action(self.robot.get_control_dict())
+            action[action_idx] = no_op_action
         return action
 
     def _reset_hand(self):
@@ -1717,11 +1709,14 @@ def _rotate_in_place(self, end_pose, angle_threshold=m.DEFAULT_ANGLE_THRESHOLD):
             base_action = action[self.robot.controller_action_idx["base"]]
 
             if not self._base_controller_is_joint:
+                base_action[0] = 0.0
                 base_action[1] = ang_vel
             else:
                 assert (
                     base_action.numel() == 3
                 ), "Currently, the action primitives only support [x, y, theta] joint controller"
+                base_action[0] = 0.0
+                base_action[1] = 0.0
                 base_action[2] = ang_vel
 
             action[self.robot.controller_action_idx["base"]] = base_action

omnigibson/controllers/controller_base.py

@@ -194,6 +194,35 @@ def _preprocess_command(self, command):
         # Return processed command
         return command
 
+    def _reverse_preprocess_command(self, processed_command):
+        """
+        Reverses the scaling operation performed by _preprocess_command.
+        Note: This method does not reverse the clipping operation as it's not reversible.
+
+        Args:
+            processed_command (th.Tensor[float]): Processed command vector
+
+        Returns:
+            th.Tensor[float]: Original command vector (before scaling, clipping not reversed)
+        """
+        # We only reverse the scaling if both input and output limits exist
+        if self._command_input_limits is not None and self._command_output_limits is not None:
+            # If we haven't calculated how to scale the command, do that now (once)
+            if self._command_scale_factor is None:
+                self._command_scale_factor = abs(self._command_output_limits[1] - self._command_output_limits[0]) / abs(
+                    self._command_input_limits[1] - self._command_input_limits[0]
+                )
+                self._command_output_transform = (self._command_output_limits[1] + self._command_output_limits[0]) / 2.0
+                self._command_input_transform = (self._command_input_limits[1] + self._command_input_limits[0]) / 2.0
+
+            original_command = (
+                processed_command - self._command_output_transform
+            ) / self._command_scale_factor + self._command_input_transform
+        else:
+            original_command = processed_command
+
+        return original_command
+
     def update_goal(self, command, control_dict):
         """
         Updates inputted @command internally, writing any necessary internal variables as needed.
@@ -305,6 +334,21 @@ def compute_no_op_goal(self, control_dict):
         """
         raise NotImplementedError
 
+    def compute_no_op_action(self, control_dict):
+        """
+        Compute no-op action given the goal
+        """
+        if self._goal is None:
+            self._goal = self.compute_no_op_goal(control_dict=control_dict)
+        command = self._compute_no_op_action(control_dict=control_dict)
+        return self._reverse_preprocess_command(command)
+
+    def _compute_no_op_action(self, control_dict):
+        """
+        Compute no-op action given the goal
+        """
+        raise NotImplementedError
+
     def _dump_state(self):
         # Default is just the command
         return dict(

omnigibson/controllers/dd_controller.py

@@ -113,6 +113,9 @@ def compute_no_op_goal(self, control_dict):
         # This is zero-vector, since we want zero linear / angular velocity
         return dict(vel=th.zeros(2))
 
+    def _compute_no_op_action(self, control_dict):
+        return self._goal["vel"]
+
     def _get_goal_shapes(self):
         # Add (2, )-array representing linear, angular velocity
         return dict(vel=(2,))

omnigibson/controllers/ik_controller.py

@@ -351,6 +351,38 @@ def compute_no_op_goal(self, control_dict):
             target_quat=control_dict[f"{self.task_name}_quat_relative"],
         )
 
+    def _compute_no_op_action(self, control_dict):
+
+        target_pos = self._goal["target_pos"]
+        target_quat = self._goal["target_quat"]
+        pos_relative = control_dict[f"{self.task_name}_pos_relative"]
+        quat_relative = control_dict[f"{self.task_name}_quat_relative"]
+
+        command = th.zeros(6, dtype=th.float32, device=target_pos.device)
+
+        # Handle position
+        if self.mode == "absolute_pose":
+            command[:3] = target_pos
+        else:
+            command[:3] = target_pos - pos_relative
+
+        # Handle orientation
+        if self.mode == "position_fixed_ori" or self.mode == "position_compliant_ori":
+            # For these modes, we don't need to add orientation to the command
+            pass
+        elif self.mode == "pose_absolute_ori" or self.mode == "absolute_pose":
+            command[3:] = T.quat2axisangle(target_quat)
+        else:  # pose_delta_ori control
+            current_rot = T.quat2mat(quat_relative)
+            target_rot = T.quat2mat(target_quat)
+            delta_rot = target_rot @ (current_rot.T)
+
+            # Convert delta rotation to axis-angle representation
+            delta_axisangle = T.quat2axisangle(T.mat2quat(delta_rot))
+            command[3:] = delta_axisangle
+
+        return command
+
     def _get_goal_shapes(self):
         return dict(
             target_pos=(3,),

omnigibson/controllers/joint_controller.py

@@ -236,6 +236,16 @@ def compute_no_op_goal(self, control_dict):
 
         return dict(target=target)
 
+    def _compute_no_op_action(self, control_dict):
+        if self._use_delta_commands:
+            base_value = control_dict[f"joint_{self._motor_type}"][self.dof_idx]
+            # TODO: do we need to care about the entire gimbal lock situation?
+            command = self._goal["target"] - base_value
+        else:
+            command = self._goal["target"]
+
+        return command
+
     def _get_goal_shapes(self):
         return dict(target=(self.control_dim,))
 

omnigibson/controllers/multi_finger_gripper_controller.py

@@ -268,6 +268,9 @@ def compute_no_op_goal(self, control_dict):
         # Just use a zero vector
         return dict(target=th.zeros(self.command_dim))
 
+    def _compute_no_op_action(self, control_dict):
+        return self._goal["target"]
+
     def _get_goal_shapes(self):
         return dict(target=(self.command_dim,))
 

omnigibson/controllers/osc_controller.py

@@ -449,6 +449,37 @@ def compute_no_op_goal(self, control_dict):
             target_ori_mat=T.quat2mat(target_quat),
         )
 
+    def _compute_no_op_action(self, control_dict):
+        target_pos = self._goal["target_pos"]
+        target_quat = T.mat2quat(self._goal["target_ori_mat"])
+        pos_relative = control_dict[f"{self.task_name}_pos_relative"]
+        quat_relative = control_dict[f"{self.task_name}_quat_relative"]
+
+        command = th.zeros(6, dtype=th.float32, device=target_pos.device)
+
+        # Handle position
+        if self.mode == "absolute_pose":
+            command[:3] = target_pos
+        else:
+            command[:3] = target_pos - pos_relative
+
+        # Handle orientation
+        if self.mode == "position_fixed_ori" or self.mode == "position_compliant_ori":
+            # For these modes, we don't need to add orientation to the command
+            pass
+        elif self.mode == "pose_absolute_ori" or self.mode == "absolute_pose":
+            command[3:] = T.quat2axisangle(target_quat)
+        else:  # pose_delta_ori control
+            current_rot = T.quat2mat(quat_relative)
+            target_rot = T.quat2mat(target_quat)
+            delta_rot = target_rot @ (current_rot.T)
+
+            # Convert delta rotation to axis-angle representation
+            delta_axisangle = T.quat2axisangle(T.mat2quat(delta_rot))
+            command[3:] = delta_axisangle
+
+        return command
+
     def _get_goal_shapes(self):
         return dict(
             target_pos=(3,),

omnigibson/maps/map_base.py

@@ -55,5 +55,5 @@ def world_to_map(self, xy):
             xy: 2D location in world reference frame (metric)
         :return: 2D location in map reference frame (image)
         """
-        point_wrt_map = th.tensor(xy) / self.map_resolution + self.map_size / 2.0
+        point_wrt_map = xy / self.map_resolution + self.map_size / 2.0
         return th.flip(point_wrt_map, dims=tuple(range(point_wrt_map.dim()))).int()

omnigibson/object_states/particle_modifier.py

@@ -26,6 +26,7 @@
     get_particle_positions_from_frame,
     get_particle_positions_in_frame,
 )
+from omnigibson.utils.numpy_utils import vtarray_to_torch
 from omnigibson.utils.python_utils import classproperty
 from omnigibson.utils.sampling_utils import sample_cuboid_on_object
 from omnigibson.utils.ui_utils import suppress_omni_log
@@ -373,7 +374,7 @@ def overlap_callback(hit):
                 if self._projection_mesh_params is None:
                     self._projection_mesh_params = {
                         "type": mesh_type,
-                        "extents": th.tensor(pre_existing_mesh.GetAttribute("xformOp:scale").Get()),
+                        "extents": vtarray_to_torch(pre_existing_mesh.GetAttribute("xformOp:scale").Get()),
                     }
                 # Otherwise, make sure we don't have a mismatch between the pre-existing shape type and the
                 # desired type since we can't delete the original mesh

omnigibson/object_states/toggle.py

@@ -8,6 +8,7 @@
 from omnigibson.object_states.update_state_mixin import GlobalUpdateStateMixin, UpdateStateMixin
 from omnigibson.prims.geom_prim import VisualGeomPrim
 from omnigibson.utils.constants import PrimType
+from omnigibson.utils.numpy_utils import vtarray_to_torch
 from omnigibson.utils.python_utils import classproperty
 from omnigibson.utils.usd_utils import RigidContactAPI, absolute_prim_path_to_scene_relative, create_primitive_mesh
 
@@ -114,7 +115,7 @@ def _initialize(self):
         else:
             # Infer radius from mesh if not specified as an input
             lazy.omni.isaac.core.utils.bounds.recompute_extents(prim=pre_existing_mesh)
-            self.scale = th.tensor(pre_existing_mesh.GetAttribute("xformOp:scale").Get())
+            self.scale = vtarray_to_torch(pre_existing_mesh.GetAttribute("xformOp:scale").Get())
 
         # Create the visual geom instance referencing the generated mesh prim
         relative_prim_path = absolute_prim_path_to_scene_relative(self.obj.scene, mesh_prim_path)

omnigibson/objects/dataset_object.py

@@ -447,7 +447,9 @@ def scales_in_link_frame(self):
                             prim.GetAttribute("physics:localRot1").Get()
                         )[[1, 2, 3, 0]]
                         # Invert the child link relationship, and multiply the two rotations together to get the final rotation
-                        local_ori = T.quat_multiply(quaternion1=T.quat_inverse(quat1), quaternion0=quat0)
+                        local_ori = T.quat_multiply(
+                            quaternion1=T.quat_inverse(th.from_numpy(quat1)), quaternion0=th.from_numpy(quat0)
+                        )
                         jnt_frame_rot = T.quat2mat(local_ori)
                         scale_in_child_lf = th.abs(jnt_frame_rot.T @ th.tensor(scale_in_parent_lf))
                         scales[child_name] = scale_in_child_lf

omnigibson/prims/cloth_prim.py

@@ -17,6 +17,7 @@
 import omnigibson.utils.transform_utils as T
 from omnigibson.macros import create_module_macros, gm
 from omnigibson.prims.geom_prim import GeomPrim
+from omnigibson.utils.numpy_utils import vtarray_to_torch
 from omnigibson.utils.sim_utils import CsRawData
 from omnigibson.utils.usd_utils import array_to_vtarray, mesh_prim_to_trimesh_mesh, sample_mesh_keypoints
 
@@ -422,7 +423,7 @@ def get_linear_velocity(self):
         Returns:
             th.tensor: current average linear velocity of the particles of the cloth prim. Shape (3,).
         """
-        return th.tensor(self._prim.GetAttribute("velocities").Get()).mean(dim=0)
+        return vtarray_to_torch(self._prim.GetAttribute("velocities").Get()).mean(dim=0)
 
     def get_angular_velocity(self):
         """

omnigibson/prims/geom_prim.py

@@ -7,6 +7,7 @@
 import omnigibson.utils.transform_utils as T
 from omnigibson.macros import gm
 from omnigibson.prims.xform_prim import XFormPrim
+from omnigibson.utils.numpy_utils import vtarray_to_torch
 from omnigibson.utils.python_utils import assert_valid_key
 from omnigibson.utils.usd_utils import PoseAPI, mesh_prim_shape_to_trimesh_mesh
 
@@ -128,7 +129,7 @@ def points(self):
         mesh_type = mesh.GetPrimTypeInfo().GetTypeName()
         if mesh_type == "Mesh":
             # If the geom is a mesh we can directly return its points.
-            return th.tensor(self.prim.GetAttribute("points").Get(), dtype=th.float32)
+            return vtarray_to_torch(self.prim.GetAttribute("points").Get(), dtype=th.float32)
         else:
             # Return the vertices of the trimesh
             return th.tensor(mesh_prim_shape_to_trimesh_mesh(mesh).vertices, dtype=th.float32)

omnigibson/prims/xform_prim.py

@@ -196,6 +196,8 @@ def set_position_orientation(self, position=None, orientation=None):
         parent_world_transform = PoseAPI.get_world_pose_with_scale(parent_path)
 
         local_transform = th.linalg.inv_ex(parent_world_transform).inverse @ my_world_transform
+        # unscale local transform's rotation
+        local_transform[:3, :3] /= th.linalg.norm(local_transform[:3, :3], dim=0)
         product = local_transform[:3, :3] @ local_transform[:3, :3].T
         assert th.allclose(
             product, th.diag(th.diag(product)), atol=1e-3

omnigibson/robots/franka.py

@@ -109,7 +109,9 @@ def __init__(
             self._finger_link_names = [f"link_{i}_0" for i in range(16)]
             self._finger_joint_names = [f"joint_{i}_0" for i in [12, 13, 14, 15, 8, 9, 10, 11, 4, 5, 6, 7, 0, 1, 2, 3]]
             # position where the hand is parallel to the ground
-            self._default_robot_model_joint_pos = th.cat(([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72], th.zeros(16)))
+            self._default_robot_model_joint_pos = th.cat(
+                (th.tensor([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72]), th.zeros(16))
+            )
             self._teleop_rotation_offset = th.tensor([0, 0.7071, 0, 0.7071])
             self._ag_start_points = [
                 GraspingPoint(link_name=f"base_link", position=[0.015, 0, -0.03]),
@@ -132,7 +134,9 @@ def __init__(
                 f"finger_joint_{i}" for i in [12, 13, 14, 15, 1, 0, 2, 3, 5, 4, 6, 7, 9, 8, 10, 11]
             ]
             # position where the hand is parallel to the ground
-            self._default_robot_model_joint_pos = th.cat(([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72], th.zeros(16)))
+            self._default_robot_model_joint_pos = th.cat(
+                (th.tensor([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72]), th.zeros(16))
+            )
             self._teleop_rotation_offset = th.tensor([-0.7071, 0.7071, 0, 0])
             self._ag_start_points = [
                 GraspingPoint(link_name=f"palm_center", position=[0, -0.025, 0.035]),
@@ -152,7 +156,9 @@ def __init__(
             self._finger_link_names = [f"link{i}" for i in hand_part_names]
             self._finger_joint_names = [f"joint{i}" for i in hand_part_names]
             # position where the hand is parallel to the ground
-            self._default_robot_model_joint_pos = th.cat(([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72], th.zeros(12)))
+            self._default_robot_model_joint_pos = th.cat(
+                (th.tensor([0.86, -0.27, -0.68, -1.52, -0.18, 1.29, 1.72]), th.zeros(12))
+            )
             self._teleop_rotation_offset = th.tensor([0, 0, 0.707, 0.707])
             # TODO: add ag support for inspire hand
             self._ag_start_points = [

omnigibson/systems/macro_particle_system.py

@@ -642,7 +642,7 @@ def generate_group_particles_on_object(self, group, max_samples=None, min_sample
                 position, normal, quaternion, hit_link, reasons = result
                 if position is not None:
                     positions.append(position)
-                    orientations.append(th.tensor(quaternion))
+                    orientations.append(quaternion)
                     particle_scales.append(scale)
                     link_prim_paths.append(hit_link)
             scales = particle_scales
@@ -1287,7 +1287,7 @@ def add_particle(self, relative_prim_path, scale, idn=None):
     def get_particles_position_orientation(self):
         # Note: This gets the center of the sphere approximation of the particles, NOT the actual particle frames!
         if self.n_particles > 0:
-            tfs = th.tensor(self.particles_view.get_transforms())
+            tfs = self.particles_view.get_transforms()
             pos, ori = tfs[:, :3], tfs[:, 3:]
             pos = pos + T.quat2mat(ori) @ self._particle_offset
         else:
@@ -1344,7 +1344,7 @@ def get_particles_velocities(self):
                 - (n, 3)-array: per-particle (ax, ay, az) angular velocities in the world frame
         """
         if self.n_particles > 0:
-            vels = th.tensor(self.particles_view.get_velocities())
+            vels = self.particles_view.get_velocities()
             lin_vel, ang_vel = vels[:, :3], vels[:, 3:]
         else:
             lin_vel, ang_vel = th.empty(0).reshape(0, 3), th.empty(0).reshape(0, 3)