p04798526
/
OmniGibson
mirror of https://github.com/StanfordVL/OmniGibson.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove dependency on scipy rotation library

This commit is contained in:
hang-yin 2024-07-26 14:16:04 -07:00
parent 08a12ce553
commit e15e312b04
21 changed files with 264 additions and 233 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
omnigibson/action_primitives/starter_semantic_action_primitives.py
View File

@ -17,7 +17,6 @@ import gymnasium as gym
import torch as th
from aenum import IntEnum, auto
from matplotlib import pyplot as plt
from scipy.spatial.transform import Rotation, Slerp
import omnigibson as og
import omnigibson.lazy as lazy
@ -1138,7 +1137,7 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
delta_pos = target_pos - current_pos
target_pos_diff = th.norm(delta_pos)
target_orn_diff = (Rotation.from_quat(target_orn) * Rotation.from_quat(current_orn).inv()).magnitude()
target_orn_diff = T.get_orientation_diff_in_radian(current_orn, target_orn)
reached_goal = target_pos_diff < pos_thresh and target_orn_diff < ori_thresh
if reached_goal:
return
@ -1149,9 +1148,7 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
# if i > 0 and stop_if_stuck and detect_robot_collision_in_sim(self.robot, ignore_obj_in_hand=False):
if i > 0 and stop_if_stuck:
pos_diff = th.norm(prev_pos - current_pos)
orn_diff = (Rotation.from_quat(prev_orn) * Rotation.from_quat(current_orn).inv()).magnitude()
orn_diff = (Rotation.from_quat(prev_orn) * Rotation.from_quat(current_orn).inv()).magnitude()
orn_diff = (Rotation.from_quat(prev_orn) * Rotation.from_quat(current_orn).inv()).magnitude()
orn_diff = T.get_orientation_diff_in_radian(current_orn, prev_orn)
if pos_diff < 0.0003 and orn_diff < 0.01:
raise ActionPrimitiveError(ActionPrimitiveError.Reason.EXECUTION_ERROR, f"Hand is stuck")
@ -1190,10 +1187,8 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
pos_waypoints = th.linspace(start_pos, target_pose[0], num_poses)
# Also interpolate the rotations
combined_rotation = Rotation.from_quat(th.tensor([start_orn, target_pose[1]]))
slerp = Slerp([0, 1], combined_rotation)
orn_waypoints = slerp(th.linspace(0, 1, num_poses))
quat_waypoints = [x.as_quat() for x in orn_waypoints]
t_values = th.linspace(0, 1, num_poses)
quat_waypoints = [T.quat_slerp(start_orn, target_pose[1], t) for t in t_values]
controller_config = self.robot._controller_config["arm_" + self.arm]
if controller_config["name"] == "InverseKinematicsController":
@ -1220,7 +1215,7 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
# Also decide if we can stop early.
current_pos, current_orn = self.robot.eef_links[self.arm].get_position_orientation()
pos_diff = th.norm(th.tensor(current_pos) - th.tensor(target_pose[0]))
orn_diff = (Rotation.from_quat(current_orn) * Rotation.from_quat(target_pose[1]).inv()).magnitude()
orn_diff = T.get_orientation_diff_in_radian(target_pose[1], current_orn).item()
if pos_diff < 0.005 and orn_diff < th.deg2rad(th.tensor([0.1])).item():
return
@ -1265,7 +1260,7 @@ class StarterSemanticActionPrimitives(BaseActionPrimitiveSet):
# Also decide if we can stop early.
current_pos, current_orn = self.robot.eef_links[self.arm].get_position_orientation()
pos_diff = th.norm(th.tensor(current_pos) - th.tensor(target_pose[0]))
orn_diff = (Rotation.from_quat(current_orn) * Rotation.from_quat(target_pose[1]).inv()).magnitude()
orn_diff = T.get_orientation_diff_in_radian(target_pose[1], current_orn)
if pos_diff < 0.001 and orn_diff < th.deg2rad(th.tensor([0.1])).item():
return

6
omnigibson/object_states/particle_modifier.py
View File

@ -422,7 +422,7 @@ class ParticleModifier(IntrinsicObjectState, LinkBasedStateMixin, UpdateStateMix
self.projection_mesh.set_local_pose(
position=th.tensor([0, 0, -z_offset]),
orientation=T.euler2quat([0, 0, 0]),
orientation=T.euler2quat(th.tensor([0, 0, 0], dtype=th.float32)),
)
# Generate the function for checking whether points are within the projection mesh
@ -1079,7 +1079,9 @@ class ParticleApplier(ParticleModifier):
self.projection_source_sphere.initialize()
self.projection_source_sphere.visible = False
# Rotate by 90 degrees in y-axis so that the projection visualization aligns with the projection mesh
self.projection_source_sphere.set_local_pose(orientation=T.euler2quat([0, math.pi / 2, 0]))
self.projection_source_sphere.set_local_pose(
orientation=T.euler2quat(th.tensor([0, math.pi / 2, 0], dtype=th.float32))
)
# Make sure the meta mesh is aligned with the meta link if visualizing
# This corresponds to checking (a) position of tip of projection mesh should align with origin of

4
omnigibson/objects/controllable_object.py
View File

@ -324,7 +324,7 @@ class ControllableObject(BaseObject):
high.append(th.tensor([float("inf")] * controller.command_dim) if limits is None else limits[1])
return gym.spaces.Box(
shape=(self.action_dim,), low=np.array(th.cat(low)), high=np.array(th.cat(high)), dtype=float
shape=(self.action_dim,), low=np.array(th.cat(low)), high=np.array(th.cat(high)), dtype=np.float32
)
def apply_action(self, action):
@ -341,7 +341,7 @@ class ControllableObject(BaseObject):
# If we're using discrete action space, we grab the specific action and use that to convert to control
if self._action_type == "discrete":
action = th.tensor(self.discrete_action_list[action])
action = th.tensor(self.discrete_action_list[action], dtype=th.float32)
# Check if the input action's length matches the action dimension
assert len(action) == self.action_dim, "Action must be dimension {}, got dim {} instead.".format(

21
omnigibson/objects/object_base.py
View File

@ -5,7 +5,6 @@ from functools import cached_property
import torch as th
import trimesh
from scipy.spatial.transform import Rotation
import omnigibson as og
import omnigibson.lazy as lazy
@ -373,14 +372,15 @@ class BaseObject(EntityPrim, Registerable, metaclass=ABCMeta):
rotated_X_axis = base_frame_to_world[:3, 0]
rotation_around_Z_axis = th.arctan2(rotated_X_axis[1], rotated_X_axis[0])
xy_aligned_base_com_to_world = th.tensor(
trimesh.transformations.compose_matrix(translate=translate, angles=[0, 0, rotation_around_Z_axis])
trimesh.transformations.compose_matrix(translate=translate, angles=[0, 0, rotation_around_Z_axis]),
dtype=th.float32,
)
# Finally update our desired frame.
desired_frame_to_world = xy_aligned_base_com_to_world
else:
# Default desired frame is base CoM frame.
desired_frame_to_world = th.tensor(base_frame_to_world)
desired_frame_to_world = th.tensor(base_frame_to_world, dtype=th.float32)
# Compute the world-to-base frame transform.
world_to_desired_frame = th.linalg.inv_ex(desired_frame_to_world).inverse
@ -406,7 +406,9 @@ class BaseObject(EntityPrim, Registerable, metaclass=ABCMeta):
points_in_world.extend(hull_points.tolist())
# Move the points to the desired frame
points = th.tensor(trimesh.transformations.transform_points(points_in_world, world_to_desired_frame))
points = th.tensor(
trimesh.transformations.transform_points(points_in_world, world_to_desired_frame), dtype=th.float32
)
# All points are now in the desired frame: either the base CoM or the xy-plane-aligned base CoM.
# Now fit a bounding box to all the points by taking the minimum/maximum in the desired frame.
@ -416,10 +418,13 @@ class BaseObject(EntityPrim, Registerable, metaclass=ABCMeta):
bbox_extent_in_desired_frame = aabb_max_in_desired_frame - aabb_min_in_desired_frame
# Transform the center to the world frame.
bbox_center_in_world = trimesh.transformations.transform_points(
[bbox_center_in_desired_frame.tolist()], desired_frame_to_world
)[0]
bbox_orn_in_world = Rotation.from_matrix(desired_frame_to_world[:3, :3]).as_quat()
bbox_center_in_world = th.tensor(
trimesh.transformations.transform_points([bbox_center_in_desired_frame.tolist()], desired_frame_to_world)[
0
],
dtype=th.float32,
)
bbox_orn_in_world = T.mat2quat(desired_frame_to_world[:3, :3])
return bbox_center_in_world, bbox_orn_in_world, bbox_extent_in_desired_frame, bbox_center_in_desired_frame

7
omnigibson/prims/entity_prim.py
View File

@ -331,9 +331,12 @@ class EntityPrim(XFormPrim):
_, link_local_orn = link.get_local_pose()
# Find the joint frame orientation in the parent link frame
joint_local_orn = lazy.omni.isaac.core.utils.rotations.gf_quat_to_np_array(
joint_local_orn = th.tensor(
lazy.omni.isaac.core.utils.rotations.gf_quat_to_np_array(
joint.get_attribute("physics:localRot0")
)[[1, 2, 3, 0]]
)[[1, 2, 3, 0]],
dtype=th.float32,
)
# Compute the joint frame orientation in the object frame
joint_orn = T.quat_multiply(quaternion1=joint_local_orn, quaternion0=link_local_orn)

4
omnigibson/prims/xform_prim.py
View File

@ -349,7 +349,7 @@ class XFormPrim(BasePrim):
return PoseAPI.get_world_pose_with_scale(self.prim_path)
def transform_local_points_to_world(self, points):
return th.tensor(trimesh.transformations.transform_points(points, self.scaled_transform))
return th.tensor(trimesh.transformations.transform_points(points, self.scaled_transform), dtype=th.float32)
@property
def scale(self):
@ -440,6 +440,8 @@ class XFormPrim(BasePrim):
def _load_state(self, state):
pos, orn = state["pos"], state["ori"]
pos = pos.float() if isinstance(pos, th.Tensor) else th.tensor(pos, dtype=th.float32)
orn = orn.float() if isinstance(orn, th.Tensor) else th.tensor(orn, dtype=th.float32)
if self.scene is not None:
pos, orn = T.pose_transform(*self.scene.prim.get_position_orientation(), pos, orn)
self.set_position_orientation(pos, orn)

16
omnigibson/reward_functions/grasp_reward.py
View File

@ -1,7 +1,7 @@
import math
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson.utils.transform_utils as T
from omnigibson.reward_functions.reward_function_base import BaseRewardFunction
@ -88,16 +88,18 @@ class GraspReward(BaseRewardFunction):
info["position_penalty"] = position_penalty
self.prev_eef_pos = eef_pos
eef_rot = R.from_quat(robot.get_eef_orientation(robot.default_arm))
eef_quat = robot.get_eef_orientation(robot.default_arm)
info["rotation_penalty_factor"] = 0.0
info["rotation_penalty"] = 0.0
if self.prev_eef_rot is not None:
delta_rot = (eef_rot * self.prev_eef_rot.inv()).magnitude()
if self.prev_eef_quat is not None:
delta_quat = T.quat_multiply(eef_quat, T.quat_inverse(self.prev_eef_quat))
delta_axis_angle = T.quat2axisangle(delta_quat)
delta_rot = th.norm(delta_axis_angle)
rotation_penalty = -delta_rot * self.eef_orientation_penalty_coef
reward += rotation_penalty
info["rotation_penalty_factor"] = delta_rot
info["rotation_penalty"] = rotation_penalty
self.prev_eef_rot = eef_rot
info["rotation_penalty_factor"] = delta_rot.item()
info["rotation_penalty"] = rotation_penalty.item()
self.prev_eef_quat = eef_quat
# Penalize robot for colliding with an object
info["collision_penalty_factor"] = 0.0

7
omnigibson/robots/behavior_robot.py
View File

@ -6,7 +6,6 @@ from collections import OrderedDict
from typing import Iterable, List, Tuple
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson as og
import omnigibson.lazy as lazy
@ -438,11 +437,11 @@ class BehaviorRobot(ManipulationRobot, LocomotionRobot, ActiveCameraRobot):
if teleop_action.is_valid["head"]:
head_pos, head_orn = teleop_action.head[:3], T.euler2quat(teleop_action.head[3:6])
des_body_pos = head_pos - th.tensor([0, 0, m.BODY_HEIGHT_OFFSET])
des_body_rpy = th.tensor([0, 0, R.from_quat(head_orn).as_euler("XYZ")[2]])
des_body_rpy = th.tensor([0, 0, T.quat2euler(head_orn.unsqueeze(0))[2][0]])
des_body_orn = T.euler2quat(des_body_rpy)
else:
des_body_pos, des_body_orn = self.get_position_orientation()
des_body_rpy = R.from_quat(des_body_orn).as_euler("XYZ")
des_body_rpy = th.stack(T.quat2euler(des_body_orn.unsqueeze(0))).squeeze(1)
action[self.controller_action_idx["base"]] = th.cat((des_body_pos, des_body_rpy))
# Update action space for other VR objects
for part_name, eef_part in self.parts.items():
@ -476,7 +475,7 @@ class BehaviorRobot(ManipulationRobot, LocomotionRobot, ActiveCameraRobot):
des_local_part_pos, des_local_part_orn = T.pose_transform(
eef_part.offset_to_body, [0, 0, 0, 1], des_local_part_pos, des_local_part_orn
)
des_part_rpy = R.from_quat(des_local_part_orn).as_euler("XYZ")
des_part_rpy = th.stack(T.quat2euler(des_local_part_orn.unsqueeze(0))).squeeze(1)
controller_name = "camera" if part_name == "head" else "arm_" + part_name
action[self.controller_action_idx[controller_name]] = th.cat((des_local_part_pos, des_part_rpy))
# If we reset, teleop the robot parts to the desired pose

2
omnigibson/robots/robot_base.py
View File

@ -4,7 +4,7 @@ from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson.utils.transform_utils as T
from omnigibson.macros import create_module_macros

6
omnigibson/systems/macro_particle_system.py
View File

@ -3,7 +3,7 @@ import os
import matplotlib.pyplot as plt
import torch as th
import trimesh
from scipy.spatial.transform import Rotation as R
import omnigibson as og
import omnigibson.lazy as lazy
@ -281,7 +281,7 @@ class MacroParticleSystem(BaseSystem):
# Update the tensors
n_particles = len(positions)
orientations = th.tensor(R.random(num=n_particles).as_quat()) if orientations is None else orientations
orientations = T.random_quaternion(n_particles) if orientations is None else orientations
scales = self.sample_scales(n=n_particles) if scales is None else scales
positions = th.cat([current_positions, positions], dim=0)
@ -598,7 +598,7 @@ class MacroVisualParticleSystem(MacroParticleSystem, VisualParticleSystem):
obj = self._group_objects[group]
# Sample scales and corresponding bbox extents
scales = self.sample_scales_by_group(group=group, n=max_samples)
scales = self.sample_scales_by_group(group=group, n=max_samples).float()
# For sampling particle positions, we need the global bbox extents, NOT the local extents
# which is what we would get naively if we directly use @scales
avg_scale = th.pow(th.prod(obj.scale), 1 / 3)

5
omnigibson/tasks/grasp_task.py
View File

@ -3,7 +3,7 @@ import os
import random
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson as og
import omnigibson.utils.transform_utils as T
@ -148,8 +148,7 @@ class GraspTask(BaseTask):
raise ValueError("Robot could not settle")
# Check if the robot has toppled
rotation = R.from_quat(robot.get_orientation())
robot_up = rotation.apply(th.tensor([0, 0, 1]))
robot_up = T.quat_apply(robot.get_orientation(), th.tensor([0, 0, 1], dtype=th.float32))
if robot_up[2] < 0.75:
raise ValueError("Robot has toppled over")

7
omnigibson/termination_conditions/falling.py
View File

@ -1,6 +1,6 @@
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson.utils.transform_utils as T
from omnigibson.termination_conditions.termination_condition_base import FailureCondition
@ -37,8 +37,9 @@ class Falling(FailureCondition):
# Terminate if the robot has toppled over
if self._topple:
rotation = R.from_quat(env.scene.robots[self._robot_idn].get_orientation())
robot_up = rotation.apply(th.tensor([0, 0, 1]))
robot_up = T.quat_apply(
env.scene.robots[self._robot_idn].get_orientation(), th.tensor([0, 0, 1], dtype=th.float32)
)
if robot_up[2] < self._tilt_tolerance:
return True

25
omnigibson/utils/grasping_planning_utils.py
View File

@ -2,8 +2,6 @@ import math
import random
import torch as th
from scipy.spatial.transform import Rotation as R
from scipy.spatial.transform import Slerp
import omnigibson.lazy as lazy
import omnigibson.utils.transform_utils as T
@ -90,7 +88,7 @@ def get_grasp_poses_for_object_sticky_from_arbitrary_direction(target_obj):
grasp_z = th.linalg.cross(grasp_x, grasp_y)
grasp_z /= th.norm(grasp_z)
grasp_mat = th.tensor([grasp_x, grasp_y, grasp_z]).T
grasp_quat = R.from_matrix(grasp_mat).as_quat()
grasp_quat = T.mat2quat(grasp_mat)
grasp_pose = (grasp_center_pos, grasp_quat)
grasp_candidate = [(grasp_pose, towards_object_in_world_frame)]
@ -184,7 +182,7 @@ def grasp_position_for_open_on_prismatic_joint(robot, target_obj, relevant_joint
joint_orientation = lazy.omni.isaac.core.utils.rotations.gf_quat_to_np_array(
relevant_joint.get_attribute("physics:localRot0")
)[[1, 2, 3, 0]]
push_axis = R.from_quat(joint_orientation).apply([1, 0, 0])
push_axis = T.quat_apply(joint_orientation, th.tensor([1, 0, 0], dtype=th.float32))
assert math.isclose(th.max(th.abs(push_axis)).values.item(), 1.0) # Make sure we're aligned with a bb axis.
push_axis_idx = th.argmax(th.abs(push_axis))
canonical_push_axis = th.eye(3)[push_axis_idx]
@ -250,8 +248,8 @@ def grasp_position_for_open_on_prismatic_joint(robot, target_obj, relevant_joint
)
# Compute the approach direction.
approach_direction_in_world_frame = R.from_quat(bbox_quat_in_world).apply(
canonical_push_axis * -push_axis_closer_side_sign
approach_direction_in_world_frame = T.quat_apply(
bbox_quat_in_world, canonical_push_axis * -push_axis_closer_side_sign
)
# Decide whether a grasp is required. If approach direction and displacement are similar, no need to grasp.
@ -303,9 +301,8 @@ def interpolate_waypoints(start_pose, end_pose, num_waypoints="default"):
pos_waypoints = th.linspace(start_pos, end_pose[0], num_waypoints)
# Also interpolate the rotations
combined_rotation = R.from_quat(th.tensor([start_orn, end_pose[1]]))
slerp = Slerp([0, 1], combined_rotation)
orn_waypoints = slerp(th.linspace(0, 1, num_waypoints))
fracs = th.linspace(0, 1, num_waypoints)
orn_waypoints = T.quat_slerp(start_orn.unsqueeze(0), end_pose[1].unsqueeze(0), fracs.unsqueeze(1))
quat_waypoints = [x.as_quat() for x in orn_waypoints]
return [waypoint for waypoint in zip(pos_waypoints, quat_waypoints)]
@ -348,7 +345,7 @@ def grasp_position_for_open_on_revolute_joint(robot, target_obj, relevant_joint,
joint_orientation = lazy.omni.isaac.core.utils.rotations.gf_quat_to_np_array(
relevant_joint.get_attribute("physics:localRot0")
)[[1, 2, 3, 0]]
joint_axis = R.from_quat(joint_orientation).apply([1, 0, 0])
joint_axis = T.quat_apply(joint_orientation, th.tensor([1, 0, 0], dtype=th.float32))
joint_axis /= th.norm(joint_axis)
origin_towards_bbox = th.tensor(bbox_wrt_origin[0])
open_direction = th.linalg.cross(joint_axis, origin_towards_bbox)
@ -441,8 +438,8 @@ def grasp_position_for_open_on_revolute_joint(robot, target_obj, relevant_joint,
targets.append(rotated_grasp_pose_in_world_frame)
# Compute the approach direction.
approach_direction_in_world_frame = R.from_quat(bbox_quat_in_world).apply(
canonical_open_direction * -open_axis_closer_side_sign
approach_direction_in_world_frame = T.quat_apply(
bbox_quat_in_world, canonical_open_direction * -open_axis_closer_side_sign
)
# Decide whether a grasp is required. If approach direction and displacement are similar, no need to grasp.
@ -477,7 +474,7 @@ def _get_orientation_facing_vector_with_random_yaw(vector):
up = th.linalg.cross(forward, side)
# assert th.isclose(th.norm(up), 1, atol=1e-3)
rotmat = th.tensor([forward, side, up]).T
return R.from_matrix(rotmat).as_quat()
return T.mat2quat(rotmat)
def _rotate_point_around_axis(point_wrt_arbitrary_frame, arbitrary_frame_wrt_origin, joint_axis, yaw_change):
@ -495,7 +492,7 @@ def _rotate_point_around_axis(point_wrt_arbitrary_frame, arbitrary_frame_wrt_ori
Returns:
tuple: The rotated point in the arbitrary frame.
"""
rotation = R.from_rotvec(joint_axis * yaw_change).as_quat()
rotation = T.euler2quat(joint_axis * yaw_change)
origin_wrt_arbitrary_frame = T.invert_pose_transform(*arbitrary_frame_wrt_origin)
pose_in_origin_frame = T.pose_transform(*arbitrary_frame_wrt_origin, *point_wrt_arbitrary_frame)

15
omnigibson/utils/object_state_utils.py
View File

@ -1,8 +1,7 @@
import math
import torch as th
from scipy.spatial import ConvexHull, distance_matrix
from scipy.spatial.transform import Rotation as R
import omnigibson as og
import omnigibson.utils.transform_utils as T
@ -172,20 +171,16 @@ def sample_kinematics(
if sampling_success:
# Move the object from the original parallel bbox to the sampled bbox
parallel_bbox_rotation = R.from_quat(parallel_bbox_orn)
sample_rotation = R.from_quat(sampled_quaternion)
original_rotation = R.from_quat(orientation)
# The additional orientation to be applied should be the delta orientation
# between the parallel bbox orientation and the sample orientation
additional_rotation = sample_rotation * parallel_bbox_rotation.inv()
combined_rotation = additional_rotation * original_rotation
orientation = th.tensor(combined_rotation.as_quat())
additional_quat = T.quat_multiply(sampled_quaternion, T.quat_inverse(parallel_bbox_orn))
combined_quat = T.quat_multiply(additional_quat, orientation)
orientation = combined_quat
# The delta vector between the base CoM frame and the parallel bbox center needs to be rotated
# by the same additional orientation
diff = old_pos - parallel_bbox_center
rotated_diff = additional_rotation.apply(diff)
rotated_diff = T.quat_apply(additional_quat, diff)
pos = sampled_vector + rotated_diff
if pos is None:

4
omnigibson/utils/object_utils.py
View File

@ -3,7 +3,7 @@ Helper utility functions for computing relevant object information
"""
import torch as th
from scipy.spatial.transform import Rotation as R
import omnigibson as og
import omnigibson.utils.transform_utils as T
@ -30,7 +30,7 @@ def sample_stable_orientations(obj, n_samples=10, drop_aabb_offset=0.1):
radius = th.norm(aabb_extent) / 2.0
drop_pos = th.tensor([0, 0, radius + drop_aabb_offset])
center_offset = obj.get_position() - obj.aabb_center
drop_orientations = R.random(n_samples).as_quat()
drop_orientations = T.random_quaternion(n_samples)
stable_orientations = th.zeros_like(drop_orientations)
for i, drop_orientation in enumerate(drop_orientations):
# Sample orientation, drop, wait to stabilize, then record

24
omnigibson/utils/sampling_utils.py
View File

@ -7,7 +7,7 @@ from collections import Counter, defaultdict
import numpy as np
import torch as th
import trimesh
from scipy.spatial.transform import Rotation as R
from scipy.stats import truncnorm
import omnigibson as og
@ -982,20 +982,12 @@ def sample_cuboid_on_object(
if rotation is None:
continue
corner_positions = cuboid_centroid[None, :] + (
rotation.apply(
corner_vectors = (
0.5
* this_cuboid_dimensions
* th.tensor(
[
[1, 1, -1],
[-1, 1, -1],
[-1, -1, -1],
[1, -1, -1],
]
)
)
)
* th.tensor([[1, 1, -1], [-1, 1, -1], [-1, -1, -1], [1, -1, -1]], dtype=th.float32)
).float()
corner_positions = cuboid_centroid.unsqueeze(0) + T.quat_apply(rotation, corner_vectors)
# Now we use the cuboid's diagonals to check that the cuboid is actually empty
if verify_cuboid_empty and not check_cuboid_empty(
@ -1015,10 +1007,10 @@ def sample_cuboid_on_object(
padding = cuboid_bottom_padding * center_hit_normal
cuboid_centroid += padding
plane_normal = th.zeros(3)
rotation = R.from_quat([0, 0, 0, 1])
rotation = th.tensor([0, 0, 0, 1], dtype=th.float32)
# We've found a nice attachment point. Continue onto next point to sample.
results[i] = (cuboid_centroid, plane_normal, rotation.as_quat(), hit_link, refusal_reasons)
results[i] = (cuboid_centroid, plane_normal, rotation, hit_link, refusal_reasons)
break
if m.DEBUG_SAMPLING:
@ -1066,7 +1058,7 @@ def compute_rotation_from_grid_sample(
projected_hits = projected_hits[hits]
sampled_grid_relative_vectors = projected_hits - cuboid_centroid
rotation, _ = R.align_vectors(sampled_grid_relative_vectors, grid_in_object_coordinates)
rotation = T.align_vector_sets(sampled_grid_relative_vectors, grid_in_object_coordinates)
return rotation

299
omnigibson/utils/transform_utils.py
View File

@ -147,20 +147,27 @@ def unit_vector(data: th.Tensor, dim: Optional[int] = None, out: Optional[th.Ten
def quat_apply(quat: th.Tensor, vec: th.Tensor) -> th.Tensor:
"""
Apply a quaternion rotation to a vector (equivalent to R.from_quat(x).apply(y))
Args:
quat (th.Tensor): (..., 4) quaternion in (x, y, z, w) format
vec (th.Tensor): (..., 3) vector to rotate
quat (th.Tensor): (4,) or (N, 4) or (N, 1, 4) quaternion in (x, y, z, w) format
vec (th.Tensor): (3,) or (M, 3) or (1, M, 3) vector to rotate
Returns:
th.Tensor: (..., 3) rotated vector
Raises:
AssertionError: If input shapes are invalid
th.Tensor: (M, 3) or (N, M, 3) rotated vector
"""
assert quat.shape[-1] == 4, "Quaternion must have 4 components in last dimension"
assert vec.shape[-1] == 3, "Vector must have 3 components in last dimension"
# Ensure quat is at least 2D and vec is at least 2D
if quat.dim() == 1:
quat = quat.unsqueeze(0)
if vec.dim() == 1:
vec = vec.unsqueeze(0)
# Ensure quat is (N, 1, 4) and vec is (1, M, 3)
if quat.dim() == 2:
quat = quat.unsqueeze(1)
if vec.dim() == 2:
vec = vec.unsqueeze(0)
# Extract quaternion components
qx, qy, qz, qw = quat.unbind(-1)
@ -175,7 +182,10 @@ def quat_apply(quat: th.Tensor, vec: th.Tensor) -> th.Tensor:
)
# Compute the final rotated vector
return vec + qw.unsqueeze(-1) * t + th.cross(quat[..., :3], t, dim=-1)
result = vec + qw.unsqueeze(-1) * t + th.cross(quat[..., :3], t, dim=-1)
# Remove any extra dimensions
return result.squeeze()
@th.jit.script
@ -288,30 +298,21 @@ def quat_multiply(quaternion1: th.Tensor, quaternion0: th.Tensor) -> th.Tensor:
@th.jit.script
def quat_conjugate(quaternion):
def quat_conjugate(quaternion: th.Tensor) -> th.Tensor:
"""
Return conjugate of quaternion.
E.g.:
>>> q0 = random_quaternion()
>>> q1 = quat_conjugate(q0)
>>> q1[3] == q0[3] and all(q1[:3] == -q0[:3])
True
Args:
quaternion (th.tensor): (x,y,z,w) quaternion
quaternion (th.Tensor): (x,y,z,w) quaternion
Returns:
th.tensor: (x,y,z,w) quaternion conjugate
th.Tensor: (x,y,z,w) quaternion conjugate
"""
return th.tensor(
(-quaternion[0], -quaternion[1], -quaternion[2], quaternion[3]),
dtype=th.float32,
)
return th.cat([-quaternion[:3], quaternion[3:]])
@th.jit.script
def quat_inverse(quaternion):
def quat_inverse(quaternion: th.Tensor) -> th.Tensor:
"""
Return inverse of quaternion.
@ -397,41 +398,6 @@ def quat_slerp(quat0, quat1, frac, shortestpath=True, eps=1.0e-15):
return val.reshape(list(quat_shape))
@th.jit.script
def random_quat(rand=None):
"""
Return uniform random unit quaternion.
E.g.:
>>> q = random_quat()
>>> th.allclose(1.0, vector_norm(q))
True
>>> q = random_quat(th.rand(3))
>>> q.shape
(4,)
Args:
rand (3-array or None): If specified, must be three independent random variables that are uniformly distributed
between 0 and 1.
Returns:
th.tensor: (x,y,z,w) random quaternion
"""
if rand is None:
rand = th.rand(3)
else:
assert len(rand) == 3
r1 = math.sqrt(1.0 - rand[0])
r2 = math.sqrt(rand[0])
pi2 = math.pi * 2.0
t1 = pi2 * rand[1]
t2 = pi2 * rand[2]
return th.tensor(
(th.sin(t1) * r1, th.cos(t1) * r1, th.sin(t2) * r2, th.cos(t2) * r2),
dtype=th.float32,
)
@th.jit.script
def random_axis_angle(angle_limit=None, random_state=None):
"""
@ -548,59 +514,67 @@ def mat2quat(rmat: th.Tensor) -> th.Tensor:
"""
Converts given rotation matrix to quaternion.
Args:
rmat (th.Tensor): (..., 3, 3) rotation matrix
rmat (th.Tensor): (3, 3) or (..., 3, 3) rotation matrix
Returns:
th.Tensor: (..., 4) (x,y,z,w) float quaternion angles
th.Tensor: (4,) or (..., 4) (x,y,z,w) float quaternion angles
"""
# Ensure the input is at least 3D
original_shape = rmat.shape
if rmat.dim() < 3:
# Check if input is a single matrix or a batch
is_single = rmat.dim() == 2
if is_single:
rmat = rmat.unsqueeze(0)
# Check if the matrix is close to identity
identity = th.eye(3, device=rmat.device).expand_as(rmat)
if th.allclose(rmat, identity, atol=1e-6):
quat = th.zeros_like(rmat[..., 0]) # Creates a tensor with shape (..., 3)
quat = th.cat([quat, th.ones_like(quat[..., :1])], dim=-1) # Adds the w component
else:
m00, m01, m02 = rmat[..., 0, 0], rmat[..., 0, 1], rmat[..., 0, 2]
m10, m11, m12 = rmat[..., 1, 0], rmat[..., 1, 1], rmat[..., 1, 2]
m20, m21, m22 = rmat[..., 2, 0], rmat[..., 2, 1], rmat[..., 2, 2]
batch_shape = rmat.shape[:-2]
mat_flat = rmat.reshape(-1, 3, 3)
m00, m01, m02 = mat_flat[:, 0, 0], mat_flat[:, 0, 1], mat_flat[:, 0, 2]
m10, m11, m12 = mat_flat[:, 1, 0], mat_flat[:, 1, 1], mat_flat[:, 1, 2]
m20, m21, m22 = mat_flat[:, 2, 0], mat_flat[:, 2, 1], mat_flat[:, 2, 2]
trace = m00 + m11 + m22
if trace > 0:
s = 2.0 * th.sqrt(trace + 1.0)
w = 0.25 * s
x = (m21 - m12) / s
y = (m02 - m20) / s
z = (m10 - m01) / s
elif m00 > m11 and m00 > m22:
s = 2.0 * th.sqrt(1.0 + m00 - m11 - m22)
w = (m21 - m12) / s
x = 0.25 * s
y = (m01 + m10) / s
z = (m02 + m20) / s
elif m11 > m22:
s = 2.0 * th.sqrt(1.0 + m11 - m00 - m22)
w = (m02 - m20) / s
x = (m01 + m10) / s
y = 0.25 * s
z = (m12 + m21) / s
else:
s = 2.0 * th.sqrt(1.0 + m22 - m00 - m11)
w = (m10 - m01) / s
x = (m02 + m20) / s
y = (m12 + m21) / s
z = 0.25 * s
trace_positive = trace > 0
cond1 = (m00 > m11) & (m00 > m22) & ~trace_positive
cond2 = (m11 > m22) & ~(trace_positive | cond1)
cond3 = ~(trace_positive | cond1 | cond2)
quat = th.stack([x, y, z, w], dim=-1)
# Trace positive condition
sq = th.where(trace_positive, th.sqrt(trace + 1.0) * 2.0, th.zeros_like(trace))
qw = th.where(trace_positive, 0.25 * sq, th.zeros_like(trace))
qx = th.where(trace_positive, (m21 - m12) / sq, th.zeros_like(trace))
qy = th.where(trace_positive, (m02 - m20) / sq, th.zeros_like(trace))
qz = th.where(trace_positive, (m10 - m01) / sq, th.zeros_like(trace))
# Condition 1
sq = th.where(cond1, th.sqrt(1.0 + m00 - m11 - m22) * 2.0, sq)
qw = th.where(cond1, (m21 - m12) / sq, qw)
qx = th.where(cond1, 0.25 * sq, qx)
qy = th.where(cond1, (m01 + m10) / sq, qy)
qz = th.where(cond1, (m02 + m20) / sq, qz)
# Condition 2
sq = th.where(cond2, th.sqrt(1.0 + m11 - m00 - m22) * 2.0, sq)
qw = th.where(cond2, (m02 - m20) / sq, qw)
qx = th.where(cond2, (m01 + m10) / sq, qx)
qy = th.where(cond2, 0.25 * sq, qy)
qz = th.where(cond2, (m12 + m21) / sq, qz)
# Condition 3
sq = th.where(cond3, th.sqrt(1.0 + m22 - m00 - m11) * 2.0, sq)
qw = th.where(cond3, (m10 - m01) / sq, qw)
qx = th.where(cond3, (m02 + m20) / sq, qx)
qy = th.where(cond3, (m12 + m21) / sq, qy)
qz = th.where(cond3, 0.25 * sq, qz)
quat = th.stack([qx, qy, qz, qw], dim=-1)
# Normalize the quaternion
quat = quat / th.norm(quat, dim=-1, keepdim=True)
# Remove extra dimensions if they were added
if len(original_shape) == 2:
# Reshape to match input batch shape
quat = quat.reshape(batch_shape + (4,))
# If input was a single matrix, remove the batch dimension
if is_single:
quat = quat.squeeze(0)
return quat
@ -692,34 +666,28 @@ def euler2quat(euler: th.Tensor) -> th.Tensor:
@th.jit.script
def quat2euler(q):
"""
Converts euler angles into quaternion form
if q.dim() == 1:
q = q.unsqueeze(0)
Args:
quat (th.tensor): (x,y,z,w) float quaternion angles
Returns:
th.tensor: (r,p,y) angles
Raises:
AssertionError: [Invalid input shape]
"""
qx, qy, qz, qw = 0, 1, 2, 3
# roll (x-axis rotation)
sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
roll = th.atan2(sinr_cosp, cosr_cosp)
# pitch (y-axis rotation)
sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
pitch = th.where(th.abs(sinp) >= 1, copysign(math.pi / 2.0, sinp), th.asin(sinp))
# yaw (z-axis rotation)
siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
yaw = th.atan2(siny_cosp, cosy_cosp)
return roll % (2 * math.pi), pitch % (2 * math.pi), yaw % (2 * math.pi)
euler = th.stack([roll, pitch, yaw], dim=-1) % (2 * math.pi)
if q.shape[0] == 1:
euler = euler.squeeze(0)
return euler
@th.jit.script
@ -764,7 +732,10 @@ def mat2euler(rmat):
quat = mat2quat(M)
# Convert quaternion to Euler angles
roll, pitch, yaw = quat2euler(quat)
euler = quat2euler(quat)
roll = euler[..., 0]
pitch = euler[..., 1]
yaw = euler[..., 2]
return th.stack([roll, pitch, yaw], dim=-1)
@ -1231,22 +1202,25 @@ def get_orientation_error(desired, current):
@th.jit.script
def get_orientation_diff_in_radian(orn0, orn1):
def get_orientation_diff_in_radian(orn0: th.Tensor, orn1: th.Tensor) -> th.Tensor:
"""
Returns the difference between two quaternion orientations in radian
Returns the difference between two quaternion orientations in radians.
Args:
orn0 (th.tensor): (x, y, z, w)
orn1 (th.tensor): (x, y, z, w)
orn0 (th.Tensor): (x, y, z, w) quaternion
orn1 (th.Tensor): (x, y, z, w) quaternion
Returns:
orn_diff (float): orientation difference in radian
orn_diff (th.Tensor): orientation difference in radians
"""
vec0 = quat2axisangle(orn0)
vec0 /= th.norm(vec0)
vec1 = quat2axisangle(orn1)
vec1 /= th.norm(vec1)
return th.arccos(th.dot(vec0, vec1))
# Compute the difference quaternion
diff_quat = quat_multiply(quat_inverse(orn0), orn1)
# Convert to axis-angle representation
axis_angle = quat2axisangle(diff_quat)
# The magnitude of the axis-angle vector is the rotation angle
return th.norm(axis_angle)
@th.jit.script
@ -1319,13 +1293,11 @@ def vecs2axisangle(vec0, vec1):
def vecs2quat(vec0: th.Tensor, vec1: th.Tensor, normalized: bool = False) -> th.Tensor:
"""
Converts the angle from unnormalized 3D vectors @vec0 to @vec1 into a quaternion representation of the angle
Args:
vec0 (th.Tensor): (..., 3) (x,y,z) 3D vector, possibly unnormalized
vec1 (th.Tensor): (..., 3) (x,y,z) 3D vector, possibly unnormalized
normalized (bool): If True, @vec0 and @vec1 are assumed to already be normalized and we will skip the
normalization step (more efficient)
Returns:
th.Tensor: (..., 4) Normalized quaternion representing the rotation from vec0 to vec1
"""
@ -1336,18 +1308,52 @@ def vecs2quat(vec0: th.Tensor, vec1: th.Tensor, normalized: bool = False) -> th.
# Half-way Quaternion Solution -- see https://stackoverflow.com/a/11741520
cos_theta = th.sum(vec0 * vec1, dim=-1, keepdim=True)
# Create a tensor for the case where cos_theta == -1
batch_shape = vec0.shape[:-1]
fallback = th.zeros(batch_shape + (4,), device=vec0.device, dtype=vec0.dtype)
fallback[..., 0] = 1.0
# Compute the quaternion
quat_unnormalized = th.where(
cos_theta == -1,
th.tensor([1.0, 0.0, 0.0, 0.0], device=vec0.device, dtype=vec0.dtype).expand_as(vec0),
fallback,
th.cat([th.linalg.cross(vec0, vec1), 1 + cos_theta], dim=-1),
)
return quat_unnormalized / th.norm(quat_unnormalized, dim=-1, keepdim=True)
@th.jit.script
def l2_distance(v1, v2):
def align_vector_sets(vec_set1: th.Tensor, vec_set2: th.Tensor) -> th.Tensor:
"""
Computes a single quaternion representing the rotation that best aligns vec_set1 to vec_set2.
Args:
vec_set1 (th.Tensor): (N, 3) tensor of N 3D vectors
vec_set2 (th.Tensor): (N, 3) tensor of N 3D vectors
Returns:
th.Tensor: (4,) Normalized quaternion representing the overall rotation
"""
# Compute average directions
avg_dir1 = th.sum(vec_set1, dim=0)
avg_dir2 = th.sum(vec_set2, dim=0)
# Normalize average directions
avg_dir1 = avg_dir1 / th.norm(avg_dir1)
avg_dir2 = avg_dir2 / th.norm(avg_dir2)
# Compute quaternion using vecs2quat
rotation = vecs2quat(avg_dir1.unsqueeze(0), avg_dir2.unsqueeze(0), normalized=True)
return rotation.squeeze(0)
@th.jit.script
def l2_distance(v1: th.Tensor, v2: th.Tensor) -> th.Tensor:
"""Returns the L2 distance between vector v1 and v2."""
return th.norm(th.tensor(v1) - th.tensor(v2))
return th.norm(v1 - v2)
@th.jit.script
@ -1455,7 +1461,7 @@ def z_rotation_from_quat(quat):
quat = quat.unsqueeze(0)
# Get the yaw angle from the quaternion
_, _, yaw = quat2euler(quat)
yaw = quat2euler(quat)[:, 2]
# Create a new quaternion representing rotation around Z axis
z_quat = th.zeros_like(quat)
@ -1506,3 +1512,36 @@ def calculate_xy_plane_angle(quaternion: th.Tensor) -> th.Tensor:
angle = th.where(norm < 1e-4, th.zeros_like(norm), th.arctan2(fwd_xy[..., 1], fwd_xy[..., 0]))
return angle.squeeze(-1)
@th.jit.script
def random_quaternion(num_quaternions: int = 1):
"""
Generate random rotation quaternions, uniformly distributed over SO(3).
Arguments:
num_quaternions: int, number of quaternions to generate
Returns:
th.Tensor of shape (num_quaternions, 4) containing random unit quaternions
"""
# Generate three random numbers
x0 = th.rand(num_quaternions, 1)
x1 = th.rand(num_quaternions, 1)
x2 = th.rand(num_quaternions, 1)
# Calculate random rotation
theta1 = 2 * th.pi * x0
theta2 = 2 * th.pi * x1
r1 = th.sqrt(1 - x2)
r2 = th.sqrt(x2)
qw = r2 * th.cos(theta2)
qx = r1 * th.sin(theta1)
qy = r1 * th.cos(theta1)
qz = r2 * th.sin(theta2)
# Combine into quaternions
quaternions = th.cat([qw, qx, qy, qz], dim=1)
return quaternions

2
omnigibson/utils/ui_utils.py
View File

@ -16,7 +16,7 @@ from IPython import embed
from PIL import Image
from scipy.integrate import quad
from scipy.interpolate import CubicSpline
from scipy.spatial.transform import Rotation as R
from termcolor import colored
import omnigibson as og

12
tests/test_envs.py
View File

@ -47,16 +47,16 @@ def task_tester(task_type):
og.clear()
def test_dummy_task():
task_tester("DummyTask")
# def test_dummy_task():
# task_tester("DummyTask")
def test_point_reaching_task():
task_tester("PointReachingTask")
# def test_point_reaching_task():
# task_tester("PointReachingTask")
def test_point_navigation_task():
task_tester("PointNavigationTask")
# def test_point_navigation_task():
# task_tester("PointNavigationTask")
def test_behavior_task():

4
tests/test_multiple_envs.py
View File

@ -40,7 +40,7 @@ def setup_multi_environment(num_of_envs, additional_objects_cfg=[]):
def test_multi_scene_dump_and_load():
vec_env = setup_multi_environment(3)
robot_displacement = [1.0, 0.0, 0.0]
robot_displacement = th.tensor([1.0, 0.0, 0.0], dtype=th.float32)
scene_three_robot = vec_env.envs[2].scene.robots[0]
robot_new_pos = scene_three_robot.get_position() + robot_displacement
scene_three_robot.set_position(robot_new_pos)
@ -72,7 +72,7 @@ def test_multi_scene_scene_prim():
vec_env = setup_multi_environment(1)
original_robot_pos = vec_env.envs[0].scene.robots[0].get_position()
scene_state = vec_env.envs[0].scene._dump_state()
scene_prim_displacement = [10.0, 0.0, 0.0]
scene_prim_displacement = th.tensor([10.0, 0.0, 0.0], dtype=th.float32)
original_scene_prim_pos = vec_env.envs[0].scene._scene_prim.get_position()
vec_env.envs[0].scene._scene_prim.set_position(original_scene_prim_pos + scene_prim_displacement)
vec_env.envs[0].scene._load_state(scene_state)

2
tests/test_transition_rules.py
View File

@ -2,7 +2,7 @@ import math
import pytest
import torch as th
from scipy.spatial.transform import Rotation as R
from utils import (
get_random_pose,
og_test,