Files
issacdataengine/workflows/simbox/solver/planner_utils.py
2026-03-16 11:44:10 +00:00

721 lines
24 KiB
Python

# pylint: skip-file
# flake8: noqa
import json
import cv2
import IPython
import numpy as np
import open3d as o3d
from colored import fg
from pydrake.all import *
from transforms3d.axangles import *
from transforms3d.euler import *
from transforms3d.quaternions import *
robot_plant = None
franka_gripper_points = np.array(
[
[0.0, 0.0, 0.0],
[0.0, 0.0, 0.0],
[0.053, -0.0, 0.075],
[-0.053, 0.0, 0.075],
[0.053, -0.0, 0.105],
[-0.053, 0.0, 0.105],
]
)
anchor_seeds = np.array(
[
[0.0, -1.285, 0, -2.356, 0.0, 1.571, 0.785, 0, 0],
[2.5, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
[2.8, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
[2, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
[2.5, 0.83, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
[0.049, 1.22, -1.87, -0.67, 2.12, 0.99, -0.85, 0, 0],
[-2.28, -0.43, 2.47, -1.35, 0.62, 2.28, -0.27, 0, 0],
[-2.02, -1.29, 2.20, -0.83, 0.22, 1.18, 0.74, 0, 0],
[-2.2, 0.03, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
[-2.5, -0.71, -2.73, -0.82, -0.7, 0.62, -0.56, 0, 0],
[-2, -0.71, -2.73, -0.82, -0.7, 0.62, -0.56, 0, 0],
[-2.66, -0.55, 2.06, -1.77, 0.96, 1.77, -1.35, 0, 0],
[1.51, -1.48, -1.12, -1.55, -1.57, 1.15, 0.24, 0, 0],
[-2.61, -0.98, 2.26, -0.85, 0.61, 1.64, 0.23, 0, 0],
]
)
def add_table_collision_free_constraint(ik, plant, frame, bb_size=[0.12, 0.08, 0.08], table_height=0.1):
# apprxoimate a link with a bounding box and add as position constraints
min_height = -0.01 + table_height
max_num = 100
y_bound = 1
ik.AddPositionConstraint(
frame,
[0, 0, 0],
plant.world_frame(),
[0, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[0, 0, -bb_size[2]],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[0, 0, bb_size[2]],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[0, -bb_size[1], 0],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[0, bb_size[1], 0],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[bb_size[0], 0, 0],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
ik.AddPositionConstraint(
frame,
[-bb_size[0], 0, 0],
plant.world_frame(),
[-0.3, -y_bound, min_height],
[max_num, y_bound, max_num],
)
def make_gripper_pts(points, color=(1, 0, 0)):
# o3d.visualization.RenderOption.line_width = 8.0
line_index = [[0, 1], [1, 2], [1, 3], [3, 5], [2, 4]]
cur_gripper_pts = points.copy()
cur_gripper_pts[1] = (cur_gripper_pts[2] + cur_gripper_pts[3]) / 2.0
line_set = o3d.geometry.LineSet()
line_set.points = o3d.utility.Vector3dVector(cur_gripper_pts)
line_set.lines = o3d.utility.Vector2iVector(line_index)
line_set.colors = o3d.utility.Vector3dVector([color for i in range(len(line_index))])
return line_set
def get_interp_time(curr_time, finish_time, ratio):
"""get interpolated time between curr and finish"""
return (finish_time - curr_time) * ratio + curr_time
def se3_inverse(RT):
RT = RT.reshape(4, 4)
R = RT[:3, :3]
T = RT[:3, 3].reshape((3, 1))
RT_new = np.eye(4, dtype=np.float32)
RT_new[:3, :3] = R.transpose()
RT_new[:3, 3] = -1 * np.dot(R.transpose(), T).reshape((3))
return RT_new
def rotAxis(axis, angle):
if axis == "x":
return rotX(angle)
elif axis == "y":
return rotY(angle)
elif axis == "z":
return rotZ(angle)
else:
Rot = np.eye(4)
Rot[:3, :3] = axangle2mat(axis, angle)
return Rot
def rotZ(rotz):
RotZ = np.array(
[
[np.cos(rotz), -np.sin(rotz), 0, 0],
[np.sin(rotz), np.cos(rotz), 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1],
]
)
return RotZ
def rotY(roty):
RotY = np.array(
[
[np.cos(roty), 0, np.sin(roty), 0],
[0, 1, 0, 0],
[-np.sin(roty), 0, np.cos(roty), 0],
[0, 0, 0, 1],
]
)
return RotY
def rotX(rotx):
RotX = np.array(
[
[1, 0, 0, 0],
[0, np.cos(rotx), -np.sin(rotx), 0],
[0, np.sin(rotx), np.cos(rotx), 0],
[0, 0, 0, 1],
]
)
return RotX
def mat2axangle_(mat):
axangle = mat2axangle(mat[:3, :3])
return axangle[0] * axangle[1]
def pack_pose(pose, rot_type="quat"): # for action
rot_index = 4 if rot_type == "quat" else 3
if rot_type == "quat":
rot_func = mat2quat
elif rot_type == "euler":
rot_func = mat2euler
elif rot_type == "axangle":
rot_func = mat2axangle_
packed = np.zeros(3 + rot_index)
packed[3 : 3 + rot_index] = rot_func(pose[:3, :3])
packed[:3] = pose[:3, 3]
return packed
def dense_sample_traj_times(sample_times, task_completion_time, start_idx=1):
"""densify the waypoints for IK trajectory"""
ik_times = list(sample_times)
for i in range(start_idx, len(sample_times)):
ratio = 1 if sample_times[i] < task_completion_time else 2
N = int((sample_times[i] - sample_times[i - 1]) * ratio) #
for j in range(N):
ik_times.insert(i, get_interp_time(sample_times[i - 1], sample_times[i], j / N))
# force insertion of sample_times and then sort/deduplicate
ik_times.extend(sample_times)
ik_times = sorted(np.unique(ik_times))
return ik_times
def compose_circular_key_frames(
center_pose,
start_time=0,
r=0.02,
discrete_frame_num=5,
total_degrees=np.pi * 2,
total_time=3,
):
"""generating the circular motion for whisk for instance"""
poses = []
times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
for i in range(discrete_frame_num):
pose = center_pose.copy()
pose[:2, 3] += (rotZ(i / discrete_frame_num * total_degrees)[:3, :3] @ np.array([r, 0, 0]))[:2]
poses.append(pose)
return times, poses
def compose_rotate_y_key_frames(
center_pose,
start_time=0,
r=0.02,
discrete_frame_num=5,
total_degrees=np.pi * 2,
total_time=3,
):
"""generating the circular motion for whisk for instance"""
poses = []
times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
for i in range(discrete_frame_num):
pose = center_pose.copy()
pose[:2, 3] += (rotZ(i / discrete_frame_num * total_degrees)[:3, :3] @ np.array([r, 0, 0]))[:2]
poses.append(pose)
return times, poses
def compose_rotating_key_frames(
ef_pose,
center_point,
start_time=0,
discrete_frame_num=5,
total_degrees=np.pi * 2,
total_time=3,
):
"""generating the circular motion with facing center for wrench for instance"""
poses = []
times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
rel_pose = ef_pose.copy()
rel_pose[:3, 3] -= center_point
for i in range(discrete_frame_num): # clockwise
pose = rotZ(-(i + 1) / discrete_frame_num * total_degrees) @ rel_pose.copy()
pose[:3, 3] += center_point
poses.append(pose)
return times, poses
def solve_ik_kpam(
link2base, # world to object link transform matrix
constraint_dicts,
gripper_frame,
keypoints_robot_in_hand, # robot keypoints in ee frame
keypoints_object_in_robot, # object keypoints in robot base frame
p0,
q0,
centering_joint,
rot_tol=np.pi / 10,
add_table_col=True,
add_gripper_faceup=False,
timeout=False,
consider_collision=False,
table_height=0.15,
contact_plane_normal=None,
):
"""
the simple case for the kpam problem
always assume the head tool keyponint match the head object keypoint
the head and tail orthogonal to the 0,0,1
and the tail and side orthogonal to 0,0,1
minimize cost and both joint space and the pose space
"""
ik_context = robot_plant.CreateDefaultContext()
ik = InverseKinematics(robot_plant, ik_context)
# separate out cost and constraint
# cost is more or less address by the position and orientation constraint
# transfer target axis orientation from link_0 frame to world frame
def get_target_axis(link2base, target_axis):
vec = link2base[:3, :3] @ target_axis
return vec
# separate out axis
for constraint in constraint_dicts:
# if "name" in constraint and constraint["name"]=="hand_parallel_to_link0_edge":
# tool_head2side = keypoints_robot_in_hand["tool_side"] - keypoints_robot_in_hand["tool_head"]
# object_vec = get_target_axis(link2base,constraint["target_axis"])
# # print("===============================",np.dot(tool_head2side, object_vec)/np.linalg.norm(tool_head2side),"===============================")
# if np.dot(tool_head2side, object_vec)/np.linalg.norm(tool_head2side)>-0.5:
# constraint["target_inner_product"]= -1
if "name" in constraint and constraint["name"] == "fingers_orthogonal_to_link0":
from_name = constraint["axis_from_keypoint_name"]
to_name = constraint["axis_to_keypoint_name"]
tol = constraint["tolerance"]
tool_vec = keypoints_robot_in_hand[to_name] - keypoints_robot_in_hand[from_name]
target1_from_name = constraint["cross_target_axis1_from_keypoint_name"]
target1_to_name = constraint["cross_target_axis1_to_keypoint_name"]
target_vec1 = keypoints_object_in_robot[target1_to_name] - keypoints_object_in_robot[target1_from_name]
target_vec2 = get_target_axis(link2base, constraint["target_axis"])
target_vec = np.cross(target_vec1, target_vec2)
tol = constraint["tolerance"]
# use cross product to get perpendicular direction
tgt = np.arccos(constraint["target_inner_product"])
lower_bound = max(tgt - tol, 0)
upper_bound = min(tgt + tol, np.pi)
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
tool_vec / np.linalg.norm(tool_vec),
robot_plant.world_frame(),
target_vec / np.linalg.norm(target_vec),
lower_bound,
upper_bound,
)
elif "name" in constraint and constraint["name"] == "hand_parallel_to_link0_edge_door":
from_name = constraint["axis_from_keypoint_name"]
to_name = constraint["axis_to_keypoint_name"]
tool_vec = keypoints_robot_in_hand[to_name] - keypoints_robot_in_hand[from_name]
target1_from_name = constraint["cross_target_axis1_from_keypoint_name"]
target1_to_name = constraint["cross_target_axis1_to_keypoint_name"]
target_vec = keypoints_object_in_robot[target1_to_name] - keypoints_object_in_robot[target1_from_name]
# use cross product to get perpendicular direction
tgt = np.arccos(constraint["target_inner_product"])
lower_bound = max(tgt - tol, 0)
upper_bound = min(tgt + tol, np.pi)
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
tool_vec / np.linalg.norm(tool_vec),
robot_plant.world_frame(),
target_vec / np.linalg.norm(target_vec),
lower_bound,
upper_bound,
)
elif "name" in constraint and constraint["name"] == "hand_parallel_to_link0_edge":
from_name = constraint["axis_from_keypoint_name"]
to_name = constraint["axis_to_keypoint_name"]
target_axis = get_target_axis(link2base, constraint["target_axis"])
vec = keypoints_robot_in_hand[to_name] - keypoints_robot_in_hand[from_name]
tol = constraint["tolerance"]
tgt_innger_product = constraint["target_inner_product"]
tgt = np.arccos(tgt_innger_product)
lower_bound = max(tgt - tol, 0)
upper_bound = min(tgt + tol, np.pi)
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
vec / np.linalg.norm(vec),
robot_plant.world_frame(),
target_axis / np.linalg.norm(target_axis),
lower_bound,
upper_bound,
)
elif "name" in constraint and constraint["name"] == "hand_parallel_to_axis_computed_by_keypoints":
from_name = constraint["axis_from_keypoint_name"]
to_name = constraint["axis_to_keypoint_name"]
vec = keypoints_robot_in_hand[to_name] - keypoints_robot_in_hand[from_name]
from_name = constraint["target_axis_from_keypoint_name"]
to_name = constraint["target_axis_to_keypoint_name"]
target_axis = keypoints_object_in_robot[to_name] - keypoints_object_in_robot[from_name]
tol = constraint["tolerance"]
tgt_innger_product = constraint["target_inner_product"]
tgt = np.arccos(tgt_innger_product)
lower_bound = max(tgt - tol, 0)
upper_bound = min(tgt + tol, np.pi)
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
vec / np.linalg.norm(vec),
robot_plant.world_frame(),
target_axis / np.linalg.norm(target_axis),
lower_bound,
upper_bound,
)
elif "name" in constraint and constraint["name"] == "hand_parallel_to_link0_move_axis":
from_name = constraint["axis_from_keypoint_name"]
to_name = constraint["axis_to_keypoint_name"]
target_axis = get_target_axis(link2base, constraint["target_axis"])
vec = keypoints_robot_in_hand[to_name] - keypoints_robot_in_hand[from_name]
tol = constraint["tolerance"]
tgt = np.arccos(constraint["target_inner_product"])
lower_bound = max(tgt - tol, 0)
upper_bound = min(tgt + tol, np.pi)
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
vec / np.linalg.norm(vec),
robot_plant.world_frame(),
target_axis / np.linalg.norm(target_axis),
lower_bound,
upper_bound,
)
elif "name" in constraint and (constraint["name"] == "fingers_contact_with_link0"):
name = constraint["keypoint_name"]
target_name = constraint["target_keypoint_name"]
tool_point = keypoints_robot_in_hand[name]
object_point = keypoints_object_in_robot[target_name]
tol = constraint["tolerance"]
ik.AddPositionConstraint(
gripper_frame,
tool_point,
robot_plant.world_frame(),
object_point - tol,
object_point + tol,
)
else:
raise ValueError("undefined constraint")
"""solving IK to match tool head keypoint and the object keypoint"""
# maybe add slackness
# make sure the arm does not go backward
ik.AddPositionConstraint(gripper_frame, [0, 0, 0], robot_plant.world_frame(), [0.05, -1, 0], [1, 1, 1])
# no rotation constraint
# ik.AddOrientationConstraint(gripper_frame, RotationMatrix(), plant.world_frame(), pose.rotation(), rot_tol)
if add_gripper_faceup:
ik.AddAngleBetweenVectorsConstraint(
gripper_frame,
[1, 0, 0],
robot_plant.world_frame(),
[0, 0, -1],
np.pi / 12,
np.pi,
)
# not touching table constraints add elbow
# if add_table_col:
# add_table_collision_free_constraint(ik, robot_plant, gripper_frame, [0.03, 0.04, 0.08])
# add_table_collision_free_constraint(
# ik,
# robot_plant,
# robot_plant.GetFrameByName("panda_link6"),
# [0.03, 0.03, 0.03],
# )
# add_table_collision_free_constraint(
# ik,
# robot_plant,
# robot_plant.GetFrameByName("panda_link7"),
# [0.03, 0.03, 0.03],
# )
if consider_collision:
ik.AddMinimumDistanceConstraint(0.01) # 0.03
prog = ik.get_mutable_prog()
q = ik.q()
solver = SnoptSolver()
# if timeout:
# solver.SetSolverOption(solver.id(), "Major Iterations Limit", 1000)
# as well as pose space costs experiments
# print("added quadratic loss")
joint_cost_mat = np.identity(len(q))
joint_cost_mat[0, 0] = 10 # 000
prog.AddQuadraticErrorCost(joint_cost_mat, centering_joint, q)
ik.AddPositionCost(gripper_frame, [0, 0, 0], robot_plant.world_frame(), p0.translation(), np.eye(3))
ik.AddOrientationCost(gripper_frame, RotationMatrix(), robot_plant.world_frame(), p0.rotation(), 1)
prog.SetInitialGuess(q, q0)
result = solver.Solve(ik.prog()) #
if result.is_success():
return result
else:
return None
def solve_ik_traj_with_standoff(
endpoint_pose,
endpoint_joints, # 9D joint vector at endpoints
q_traj, # trajectory interpolated from current -> target -> current pose
waypoint_times, # time stamps after interpolation densification
keytimes, # keyframe time stamps
keyposes, # keyframe poses as transform matrices
rot_tol=0.03,
):
"""run a small trajopt on the trajectory with the solved IK from end-effector traj"""
# make sure the beginning and the end do not get updated
# print(waypoint_times)
# print(keytimes)
# print(len(keyposes))
waypoint_num = len(waypoint_times)
prog = MathematicalProgram()
q = prog.NewContinuousVariables(9, waypoint_num)
gripper_frame = robot_plant.GetFrameByName("panda_hand")
plant_contexts = [robot_plant.CreateDefaultContext() for i in range(waypoint_num)]
q0 = np.array([q_traj.value(t) for t in waypoint_times])
# Add key poses as constraints into the optimization program
for idx, time in enumerate(waypoint_times):
if time in keytimes: # standoff
keypoint_idx = keytimes.index(time)
if keypoint_idx == 0:
continue
pose = RigidTransform(keyposes[keypoint_idx])
prog.AddConstraint(
OrientationConstraint(
robot_plant,
gripper_frame,
RotationMatrix(),
robot_plant.world_frame(),
pose.rotation(),
rot_tol,
plant_contexts[idx],
),
q[:, idx],
)
prog.AddConstraint(
PositionConstraint(
robot_plant,
robot_plant.world_frame(),
pose.translation(),
pose.translation(),
gripper_frame,
[0, 0, 0],
plant_contexts[idx],
),
q[:, idx],
)
# matrix=np.eye(4)
# matrix[:3,:3]=pose.rotation()
# matrix[:3,3]=pose.translation()
# print(pose.GetAsMatrix4()[0],",")
# # table constraint
# prog.AddConstraint(
# PositionConstraint(
# robot_plant,
# robot_plant.world_frame(),
# [0.05, -0.7, 0.06],
# [1, 0.7, 1],
# gripper_frame,
# [0, 0, 0],
# plant_contexts[idx],
# ),
# q[:, idx],
# )
# Constrain the initial trajectory point to match the current pose
# add some other constraints
prog.AddConstraint(np.sum((q[:, 0] - endpoint_joints[:, 0]) ** 2) == 0) # same starting point
# Constrain trajectory smoothness
# Add smoothness cost
weight = np.ones((9, 1))
weight[0] = 10.0
weight[-1] = 10.0
prog.AddQuadraticCost(np.sum(weight * (q[:, 1:] - q[:, :-1]) ** 2))
prog.SetInitialGuess(q, q0.squeeze().T)
# solve the trajopt
try:
solver = SnoptSolver()
result = solver.Solve(prog)
if not result.is_success():
print("SnoptSolver failed, use IpoptSolver")
solver = IpoptSolver()
result = solver.Solve(prog)
except:
print("SnoptSolver error, use IpoptSolver")
solver = IpoptSolver()
result = solver.Solve(prog)
if result.is_success():
return result
else:
print("solve_ik_traj_with_standoff return None")
return None
def build_plant(time_step=0.004, discrete_contact_solver="sap", robot_name=None, init_qpos=None):
global robot_plant
multibody_plant_config = MultibodyPlantConfig(
time_step=time_step,
discrete_contact_solver=discrete_contact_solver,
)
builder = DiagramBuilder()
robot_plant, _ = AddMultibodyPlant(multibody_plant_config, builder)
if "split_aloha" in robot_name:
# franka = AddFranka(robot_plant, init_qpos=None)
franka = AddPiper(robot_plant, init_qpos=None)
elif "lift2" in robot_name:
franka = AddR5a(robot_plant, init_qpos=None)
elif "franka" in robot_name:
# if "robotiq" in robot_name:
# franka = AddFranka(robot_plant, init_qpos=[0, -np.pi / 4, 0, -3 * np.pi / 4, 0, np.pi / 2, np.pi / 4])
# else:
# franka = AddFranka(robot_plant, init_qpos=[0, -np.pi / 4, 0, -3 * np.pi / 4, 0, np.pi / 2, np.pi / 4])
# set_trace()
franka = AddFranka(robot_plant, init_qpos=init_qpos)
else:
raise NotImplementedError
robot_plant.Finalize()
diagram = builder.Build()
fk_context = diagram.CreateDefaultContext()
fk_plant_context = robot_plant.GetMyMutableContextFromRoot(fk_context)
return robot_plant, fk_plant_context
def AddR5a(plant, init_qpos):
franka_combined_path = "workflows/simbox/panda_drake/r5a/R5a.urdf"
parser = Parser(plant)
franka = parser.AddModelFromFile(franka_combined_path)
plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base_link"))
# Set default positions:
if init_qpos is None:
init_qpos = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.04, 0.04]
index = 0
for joint_index in plant.GetJointIndices(franka):
joint = plant.get_mutable_joint(joint_index)
if isinstance(joint, RevoluteJoint):
joint.set_default_angle(init_qpos[index])
index += 1
return franka
def AddPiper(plant, init_qpos):
franka_combined_path = "workflows/simbox/panda_drake/piper100/piper100.urdf"
parser = Parser(plant)
franka = parser.AddModelFromFile(franka_combined_path)
plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("arm_base"))
# Set default positions:
if init_qpos is None:
init_qpos = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.05, -0.05]
index = 0
for joint_index in plant.GetJointIndices(franka):
joint = plant.get_mutable_joint(joint_index)
if isinstance(joint, RevoluteJoint):
joint.set_default_angle(init_qpos[index])
index += 1
return franka
def AddFranka(plant, init_qpos):
franka_combined_path = "workflows/simbox/panda_drake/panda_arm_hand.urdf"
parser = Parser(plant)
franka = parser.AddModelFromFile(franka_combined_path)
plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("panda_link0"))
# Set default positions:
if init_qpos is None:
init_qpos = [
# 0,
# -0.32,
# 0.0,
# -2.617993877991494,
# 0.0,
# 2.23,
# 0.7853981633974483,
# 0.04,
# 0.04,
0,
-math.pi / 4,
0,
-3 * math.pi / 4,
0,
math.pi / 2,
0.0,
0.0,
0.0,
]
index = 0
for joint_index in plant.GetJointIndices(franka):
joint = plant.get_mutable_joint(joint_index)
if isinstance(joint, RevoluteJoint):
joint.set_default_angle(init_qpos[index])
index += 1
return franka