Fix type errors in og-develop.

This commit is contained in:
Albert Yu 2024-09-24 02:21:58 -05:00
parent 6d3c4923f6
commit dfc6623d44
12 changed files with 45 additions and 38 deletions

View File

@ -187,6 +187,7 @@ class PlanningContext(object):
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 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
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 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
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 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
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),

View File

@ -133,8 +133,8 @@ class BaseController(Serializable, Registerable, Recreatable):
)
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 @@ class BaseController(Serializable, Registerable, Recreatable):
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
)

View File

@ -312,7 +312,7 @@ class InverseKinematicsController(JointController, ManipulationController):
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 @@ class InverseKinematicsController(JointController, ManipulationController):
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

View File

@ -135,7 +135,7 @@ class JointController(LocomotionController, ManipulationController, GripperContr
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 @@ class JointController(LocomotionController, ManipulationController, GripperContr
# 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,7 +189,7 @@ class JointController(LocomotionController, ManipulationController, GripperContr
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
@ -197,7 +197,7 @@ class JointController(LocomotionController, ManipulationController, GripperContr
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 @@ class JointController(LocomotionController, ManipulationController, GripperContr
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 @@ class JointController(LocomotionController, ManipulationController, GripperContr
# 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)

View File

@ -159,20 +159,20 @@ class MultiFingerGripperController(GripperController):
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 @@ class MultiFingerGripperController(GripperController):
# 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 @@ class MultiFingerGripperController(GripperController):
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 @@ class MultiFingerGripperController(GripperController):
# 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)

View File

@ -436,7 +436,7 @@ class ControllableObject(BaseObject):
# 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"]

View File

@ -1018,7 +1018,7 @@ class EntityPrim(XFormPrim):
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:

View File

@ -344,7 +344,7 @@ class XFormPrim(BasePrim):
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 @@ class XFormPrim(BasePrim):
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):
"""

View File

@ -1306,7 +1306,7 @@ class ManipulationRobot(BaseRobot):
# 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,

View File

@ -529,9 +529,10 @@ class ArticulationView(_ArticulationView):
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,

View File

@ -106,6 +106,7 @@ def plan_base_motion(
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 plan_base_motion(
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 plan_arm_motion_ik(
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)

View File

@ -129,7 +129,7 @@ class Remapper:
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 = {}