Fix tensor type issues on og-develop

omnigibson/action_primitives/starter_semantic_action_primitives.py

@@ -187,6 +187,7 @@ def _assemble_robot_copy(self):
                     if link_name in arm_links
                     else self.robot_copy.links_relative_poses[self.robot_copy_type][link_name]
                 )
+                link_pose = [th.tensor(arr) for arr in link_pose]
                 mesh_copy_pose = T.pose_transform(
                     *link_pose, *self.robot_copy.relative_poses[self.robot_copy_type][link_name][mesh_name]
                 )
@@ -1040,7 +1041,7 @@ def _move_hand_ik(self, eef_pose, stop_if_stuck=False):
         indented_print("Plan has %d steps", len(plan))
         for i, target_pose in enumerate(plan):
             target_pos = target_pose[:3]
-            target_quat = T.axisangle2quat(target_pose[3:])
+            target_quat = T.axisangle2quat(th.tensor(target_pose[3:]))
             indented_print("Executing grasp plan step %d/%d", i + 1, len(plan))
             yield from self._move_hand_direct_ik(
                 (target_pos, target_quat), ignore_failure=True, in_world_frame=False, stop_if_stuck=stop_if_stuck
@@ -1146,6 +1147,8 @@ def _move_hand_direct_ik(
         assert controller_config["mode"] == "pose_absolute_ori", "Controller must be in pose_absolute_ori mode"
         if in_world_frame:
             target_pose = self._get_pose_in_robot_frame(target_pose)
+        else:
+            target_pose = [th.tensor(arr) for arr in target_pose]
         target_pos = target_pose[0]
         target_orn = target_pose[1]
         target_orn_axisangle = T.quat2axisangle(target_pose[1])
@@ -1759,8 +1762,8 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs):
                 distance_lo, distance_hi = 0.0, 5.0
                 distance = (th.rand(1) * (distance_hi - distance_lo) + distance_lo).item()
                 yaw_lo, yaw_hi = -math.pi, math.pi
-                yaw = (th.rand(1) * (yaw_hi - yaw_lo) + yaw_lo).item()
-                avg_arm_workspace_range = th.mean(self.robot.arm_workspace_range[self.arm])
+                yaw = th.rand(1) * (yaw_hi - yaw_lo) + yaw_lo
+                avg_arm_workspace_range = th.mean(th.tensor(self.robot.arm_workspace_range[self.arm]))
                 pose_2d = th.tensor(
                     [
                         pose_on_obj[0][0] + distance * th.cos(yaw),

omnigibson/controllers/controller_base.py

@@ -133,8 +133,8 @@ def __init__(
         )
         command_output_limits = (
             (
-                self._control_limits[self.control_type][0][self.dof_idx],
-                self._control_limits[self.control_type][1][self.dof_idx],
+                th.tensor(self._control_limits[self.control_type][0])[self.dof_idx.long()],
+                th.tensor(self._control_limits[self.control_type][1])[self.dof_idx.long()],
             )
             if type(command_output_limits) == str and command_output_limits == "default"
             else command_output_limits
@@ -281,11 +281,11 @@ def clip_control(self, control):
             Array[float]: Clipped control signal
         """
         clipped_control = control.clip(
-            self._control_limits[self.control_type][0][self.dof_idx],
-            self._control_limits[self.control_type][1][self.dof_idx],
+            self._control_limits[self.control_type][0][self.dof_idx.long()],
+            self._control_limits[self.control_type][1][self.dof_idx.long()],
         )
         idx = (
-            self._dof_has_limits[self.dof_idx]
+            self._dof_has_limits[self.dof_idx.long()]
             if self.control_type == ControlType.POSITION
             else [True] * self.control_dim
         )

omnigibson/controllers/ik_controller.py

@@ -312,7 +312,7 @@ def compute_control(self, goal_dict, control_dict):
         target_quat = goal_dict["target_quat"]
 
         # Calculate and return IK-backed out joint angles
-        current_joint_pos = control_dict["joint_position"][self.dof_idx]
+        current_joint_pos = control_dict["joint_position"][self.dof_idx.long()]
 
         # 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.
@@ -326,15 +326,15 @@ def compute_control(self, goal_dict, control_dict):
             err = th.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 = control_dict[f"{self.task_name}_jacobian_relative"][:, self.dof_idx.long()]
             j_eef_pinv = th.linalg.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.clamp(
-                min=self._control_limits[ControlType.get_type("position")][0][self.dof_idx],
-                max=self._control_limits[ControlType.get_type("position")][1][self.dof_idx],
+                min=self._control_limits[ControlType.get_type("position")][0][self.dof_idx.long()],
+                max=self._control_limits[ControlType.get_type("position")][1][self.dof_idx.long()],
             )
 
         # Optionally pass through smoothing filter for better stability

omnigibson/controllers/joint_controller.py

@@ -135,7 +135,7 @@ def __init__(
 
     def _update_goal(self, command, control_dict):
         # Compute the base value for the command
-        base_value = control_dict[f"joint_{self._motor_type}"][self.dof_idx]
+        base_value = control_dict[f"joint_{self._motor_type}"][self.dof_idx.long()]
 
         # If we're using delta commands, add this value
         if self._use_delta_commands:
@@ -166,8 +166,8 @@ def _update_goal(self, command, control_dict):
 
         # Clip the command based on the limits
         target = target.clip(
-            self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx],
-            self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx],
+            self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx.long()],
+            self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx.long()],
         )
 
         return dict(target=target)
@@ -189,15 +189,15 @@ def compute_control(self, goal_dict, control_dict):
         Returns:
             Array[float]: outputted (non-clipped!) control signal to deploy
         """
-        base_value = control_dict[f"joint_{self._motor_type}"][self.dof_idx]
+        base_value = control_dict[f"joint_{self._motor_type}"][self.dof_idx.long()]
         target = goal_dict["target"]
 
         # Convert control into efforts
         if self._use_impedances:
             if self._motor_type == "position":
                 # Run impedance controller -- effort = pos_err * kp + vel_err * kd
                 position_error = target - base_value
-                vel_pos_error = -control_dict[f"joint_velocity"][self.dof_idx]
+                vel_pos_error = -control_dict[f"joint_velocity"][self.dof_idx.long()]
                 u = position_error * self.kp + vel_pos_error * self.kd
             elif self._motor_type == "velocity":
                 # Compute command torques via PI velocity controller plus gravity compensation torques
@@ -207,16 +207,17 @@ def compute_control(self, goal_dict, control_dict):
                 u = target
 
             dof_idxs_mat = th.meshgrid(self.dof_idx, self.dof_idx, indexing="xy")
+            dof_idxs_mat = tuple(x.long() for x in dof_idxs_mat)
             mm = control_dict["mass_matrix"][dof_idxs_mat]
             u = mm @ u
 
             # Add gravity compensation
             if self._use_gravity_compensation:
-                u += control_dict["gravity_force"][self.dof_idx]
+                u += control_dict["gravity_force"][self.dof_idx.long()]
 
             # Add Coriolis / centrifugal compensation
             if self._use_cc_compensation:
-                u += control_dict["cc_force"][self.dof_idx]
+                u += control_dict["cc_force"][self.dof_idx.long()]
 
         else:
             # Desired is the exact goal
@@ -229,7 +230,7 @@ def compute_no_op_goal(self, control_dict):
         # Compute based on mode
         if self._motor_type == "position":
             # Maintain current qpos
-            target = control_dict[f"joint_{self._motor_type}"][self.dof_idx]
+            target = control_dict[f"joint_{self._motor_type}"][self.dof_idx.long()]
         else:
             # For velocity / effort, directly set to 0
             target = th.zeros(self.control_dim)

omnigibson/controllers/multi_finger_gripper_controller.py

@@ -159,20 +159,20 @@ def compute_control(self, goal_dict, control_dict):
             Array[float]: outputted (non-clipped!) control signal to deploy
         """
         target = goal_dict["target"]
-        joint_pos = control_dict["joint_position"][self.dof_idx]
+        joint_pos = control_dict["joint_position"][self.dof_idx.long()]
         # Choose what to do based on control mode
         if self._mode == "binary":
             # Use max control signal
             should_open = target[0] >= 0.0 if not self._inverted else target[0] > 0.0
             if should_open:
                 u = (
-                    self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx]
+                    self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx.long()]
                     if self._open_qpos is None
                     else self._open_qpos
                 )
             else:
                 u = (
-                    self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx]
+                    self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx.long()]
                     if self._closed_qpos is None
                     else self._closed_qpos
                 )
@@ -183,10 +183,10 @@ def compute_control(self, goal_dict, control_dict):
         # If we're near the joint limits and we're using velocity / torque control, we zero out the action
         if self._motor_type in {"velocity", "torque"}:
             violate_upper_limit = (
-                joint_pos > self._control_limits[ControlType.POSITION][1][self.dof_idx] - self._limit_tolerance
+                joint_pos > self._control_limits[ControlType.POSITION][1][self.dof_idx.long()] - self._limit_tolerance
             )
             violate_lower_limit = (
-                joint_pos < self._control_limits[ControlType.POSITION][0][self.dof_idx] + self._limit_tolerance
+                joint_pos < self._control_limits[ControlType.POSITION][0][self.dof_idx.long()] + self._limit_tolerance
             )
             violation = th.logical_or(violate_upper_limit * (u > 0), violate_lower_limit * (u < 0))
             u *= ~violation
@@ -228,7 +228,7 @@ def _update_grasping_state(self, control_dict):
             is_grasping = IsGraspingState.UNKNOWN
 
         else:
-            finger_pos = control_dict["joint_position"][self.dof_idx]
+            finger_pos = control_dict["joint_position"][self.dof_idx.long()]
 
             # For joint position control, if the desired positions are the same as the current positions, is_grasping unknown
             if self._motor_type == "position" and th.mean(th.abs(finger_pos - self._control)) < m.POS_TOLERANCE:
@@ -240,9 +240,9 @@ def _update_grasping_state(self, control_dict):
 
             # Otherwise, the last control signal intends to "move" the gripper
             else:
-                finger_vel = control_dict["joint_velocity"][self.dof_idx]
-                min_pos = self._control_limits[ControlType.POSITION][0][self.dof_idx]
-                max_pos = self._control_limits[ControlType.POSITION][1][self.dof_idx]
+                finger_vel = control_dict["joint_velocity"][self.dof_idx.long()]
+                min_pos = self._control_limits[ControlType.POSITION][0][self.dof_idx.long()]
+                max_pos = self._control_limits[ControlType.POSITION][1][self.dof_idx.long()]
 
                 # Make sure we don't have any invalid values (i.e.: fingers should be within the limits)
                 finger_pos = th.clip(finger_pos, min_pos, max_pos)

omnigibson/objects/controllable_object.py

@@ -436,7 +436,7 @@ def step(self):
         # By default, the control type is None and the control value is 0 (th.zeros) - i.e. no control applied
         u_type_vec = th.tensor([ControlType.NONE] * self.n_dof)
         for group, ctrl in control.items():
-            idx = self._controllers[group].dof_idx
+            idx = self._controllers[group].dof_idx.long()
             u_vec[idx] = ctrl["value"]
             u_type_vec[idx] = ctrl["type"]
 

omnigibson/prims/entity_prim.py

@@ -1018,7 +1018,7 @@ def set_position_orientation(self, position=None, orientation=None, frame: Liter
             frame (Literal): The frame in which to set the position and orientation. Defaults to world.
                 scene frame sets position relative to the scene.
         """
-        assert frame in ["world", "scene"], f"Invalid frame '{frame}'. Must be 'world', or 'scene'."
+        assert frame in ["world", "scene", "parent"], f"Invalid frame '{frame}'. Must be 'world', 'scene', or 'parent'"
 
         # If kinematic only, clear cache for the root link
         if self.kinematic_only:

omnigibson/prims/xform_prim.py

@@ -344,7 +344,7 @@ def get_local_pose(self):
         logger.warning(
             'get_local_pose is deprecated and will be removed in a future release. Use get_position_orientation(frame="parent") instead'
         )
-        return self.get_position_orientation(self.prim_path, frame="parent")
+        return self.get_position_orientation(frame="parent")
 
     def set_local_pose(self, position=None, orientation=None):
         """
@@ -359,7 +359,7 @@ def set_local_pose(self, position=None, orientation=None):
         logger.warning(
             'set_local_pose is deprecated and will be removed in a future release. Use set_position_orientation(position=position, orientation=orientation, frame="parent") instead'
         )
-        return self.set_position_orientation(self.prim_path, position, orientation, frame="parent")
+        return self.set_position_orientation(position, orientation, frame="parent")
 
     def get_world_scale(self):
         """

omnigibson/robots/manipulation_robot.py

@@ -1306,7 +1306,7 @@ def _handle_assisted_grasping(self):
             # stays the same across different controllers and control modes (absolute / delta). This way,
             # a zero action will actually keep the AG setting where it already is.
             controller = self._controllers[f"gripper_{arm}"]
-            controlled_joints = controller.dof_idx
+            controlled_joints = controller.dof_idx.long()
             threshold = th.mean(
                 th.stack([self.joint_lower_limits[controlled_joints], self.joint_upper_limits[controlled_joints]]),
                 dim=0,

omnigibson/utils/deprecated_utils.py

@@ -529,9 +529,10 @@ def set_joint_positions(
             carb.log_warn("ArticulationView needs to be initialized.")
             return
         if not omni.timeline.get_timeline_interface().is_stopped() and self._physics_view is not None:
+            positions = positions.float()
             indices = self._backend_utils.resolve_indices(indices, self.count, self._device)
             joint_indices = self._backend_utils.resolve_indices(joint_indices, self.num_dof, self._device)
-            new_dof_pos = self._physics_view.get_dof_positions()
+            new_dof_pos = self._physics_view.get_dof_positions().float()
             new_dof_pos = self._backend_utils.assign(
                 self._backend_utils.move_data(positions, device=self._device),
                 new_dof_pos,

omnigibson/utils/motion_planning_utils.py

@@ -106,6 +106,7 @@ def get_angle_between_poses(p1, p2):
             segment = []
             segment.append(p2[0] - p1[0])
             segment.append(p2[1] - p1[1])
+            segment = th.tensor(segment)
             return th.arctan2(segment[1], segment[0])
 
     def create_state(space, x, y, yaw):
@@ -122,7 +123,7 @@ def state_valid_fn(q):
         x = q.getX()
         y = q.getY()
         yaw = q.getYaw()
-        pose = ([x, y, 0.0], T.euler2quat((0, 0, yaw)))
+        pose = ([x, y, 0.0], T.euler2quat(th.tensor([0, 0, yaw])))
         return not set_base_and_detect_collision(context, pose)
 
     def remove_unnecessary_rotations(path):
@@ -364,7 +365,7 @@ def state_valid_fn(q):
         eef_pose = [q[i] for i in range(6)]
         control_joint_pos = ik_solver.solve(
             target_pos=eef_pose[:3],
-            target_quat=T.axisangle2quat(eef_pose[3:]),
+            target_quat=T.axisangle2quat(th.tensor(eef_pose[3:])),
             max_iterations=1000,
         )
 
@@ -479,6 +480,7 @@ def set_arm_and_detect_collision(context, joint_pos):
         if link in robot_copy.meshes[robot_copy_type].keys():
             for mesh_name, mesh in robot_copy.meshes[robot_copy_type][link].items():
                 relative_pose = robot_copy.relative_poses[robot_copy_type][link][mesh_name]
+                pose = [th.tensor(arr) for arr in pose]
                 mesh_pose = T.pose_transform(*pose, *relative_pose)
                 translation = lazy.pxr.Gf.Vec3d(*th.tensor(mesh_pose[0], dtype=th.float32).tolist())
                 mesh.GetAttribute("xformOp:translate").Set(translation)

omnigibson/utils/vision_utils.py

@@ -129,7 +129,7 @@ def remap(self, old_mapping, new_mapping, image, image_keys=None):
                 self.key_array[key] = new_key
 
         # Apply remapping
-        remapped_img = self.key_array[image]
+        remapped_img = self.key_array[image.long()]
         # Make sure all values are correctly remapped and not equal to the default value
         assert th.all(remapped_img != th.iinfo(th.int32).max), "Not all keys in the image are in the key array!"
         remapped_labels = {}