167 lines
6.2 KiB
Python
167 lines
6.2 KiB
Python
import numpy as np
|
|
from scipy.interpolate import RegularGridInterpolator
|
|
from scipy.optimize import dual_annealing
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
from task_gen_dependencies.cost_function import get_cost_func
|
|
from task_gen_dependencies.transform_utils import unnormalize_vars, \
|
|
normalize_vars, pose2mat, euler2quat
|
|
|
|
np.random.seed(0)
|
|
|
|
|
|
# align
|
|
def generate_random_pose(mesh, plane_size):
|
|
bounding_box = mesh.bounds
|
|
min_corner = bounding_box[0]
|
|
max_corner = bounding_box[1]
|
|
|
|
object_width = max_corner[0] - min_corner[0]
|
|
object_height = max_corner[1] - min_corner[1]
|
|
position_range_x = (plane_size[0] / 2) - (object_width / 2)
|
|
position_range_y = (plane_size[1] / 2) - (object_height / 2)
|
|
position_x = np.random.uniform(-position_range_x, position_range_x)
|
|
position_y = np.random.uniform(-position_range_y, position_range_y)
|
|
yaw = np.random.uniform(0, 360)
|
|
|
|
rotation = R.from_euler('z', yaw, degrees=True).as_matrix()
|
|
pose = np.eye(4)
|
|
pose[:3, :3] = rotation
|
|
pose[:3, 3] = [position_x, position_y, -min_corner[2] + 0.002]
|
|
|
|
return pose
|
|
|
|
|
|
def passive_setup_sdf(bounding_box_range, sdf_voxels):
|
|
# create callable sdf function with interpolation
|
|
min_corner = bounding_box_range[0]
|
|
max_corner = bounding_box_range[1]
|
|
|
|
x = np.linspace(min_corner[0], max_corner[0], sdf_voxels.shape[0])
|
|
y = np.linspace(min_corner[1], max_corner[1], sdf_voxels.shape[1])
|
|
z = np.linspace(min_corner[2], max_corner[2], sdf_voxels.shape[2])
|
|
sdf_func = RegularGridInterpolator((x, y, z), sdf_voxels, bounds_error=False, fill_value=0)
|
|
return sdf_func
|
|
|
|
|
|
def passive_setup_sdf_all(bounds_max, bounds_min, sdf_voxels):
|
|
x = np.linspace(bounds_min[0], bounds_max[0], sdf_voxels.shape[0])
|
|
y = np.linspace(bounds_min[1], bounds_max[1], sdf_voxels.shape[1])
|
|
z = np.linspace(bounds_min[2], bounds_max[2], sdf_voxels.shape[2])
|
|
sdf_func = RegularGridInterpolator((x, y, z), sdf_voxels, bounds_error=False, fill_value=0)
|
|
return sdf_func
|
|
|
|
|
|
|
|
|
|
|
|
class LayoutSolver3D:
|
|
def __init__(self, workspace_xyz, workspace_size, objects=None, obj_infos=None):
|
|
x_half, y_half, z_half = workspace_size / 2
|
|
bounds = {"min_bound": [-x_half, -y_half, 0.1], "max_bound": [x_half, y_half, z_half]}
|
|
self.bounds_min = bounds["min_bound"]
|
|
self.bounds_max = bounds["max_bound"]
|
|
self.workspace_size = np.array(bounds["max_bound"]) - np.array(bounds["min_bound"])
|
|
|
|
self.objects = objects
|
|
|
|
def init_pose(self, opt_obj_ids, constraint):
|
|
active_obj_id, passive_obj_id, constraint = constraint['active'], constraint['passive'], constraint[
|
|
'constraint']
|
|
obj_active = self.objects[active_obj_id]
|
|
obj_passive = self.objects[passive_obj_id]
|
|
|
|
# if 'out' in constraints or 'on' in constraints or 'in' in constraints:
|
|
# pose_passive = generate_random_pose(obj_passive.mesh, self.workspace_size)
|
|
# self.update_obj(passive_obj_id, pose_passive)
|
|
|
|
pose_active = generate_random_pose(obj_active.mesh, self.workspace_size)
|
|
self.update_obj(active_obj_id, pose_active)
|
|
|
|
self.safe_distance = np.mean(obj_active.size[:2]) + np.mean(obj_passive.size[:2]) / 2 + 5
|
|
|
|
optimize_z = constraint == 'on'
|
|
if optimize_z:
|
|
pose_bounds_min_ac = np.append(self.bounds_min[:2], obj_passive.obj_pose[2, 3] + obj_passive.size[2] / 2.0)
|
|
pose_bounds_max_ac = np.append(self.bounds_max[:2], self.bounds_max[2])
|
|
else:
|
|
pose_bounds_min_ac = np.append(self.bounds_min[:2], pose_active[2, 3])
|
|
pose_bounds_max_ac = np.append(self.bounds_max[:2], pose_active[2, 3])
|
|
|
|
# Define bounds for optimization
|
|
rot_bounds_min = np.array([0])
|
|
rot_bounds_max = np.array([np.radians(360)])
|
|
self.og_bounds = [(b_min, b_max) for b_min, b_max in zip(
|
|
np.concatenate([pose_bounds_min_ac, rot_bounds_min]),
|
|
np.concatenate([pose_bounds_max_ac, rot_bounds_max])
|
|
)]
|
|
self.norm_bounds = np.array([(-1, 1)] * len(self.og_bounds))
|
|
|
|
# Initialize optimization
|
|
translation_active = pose_active[:3, 3]
|
|
rotation_active = R.from_matrix(pose_active[:3, :3])
|
|
euler_angles_active = rotation_active.as_euler('xyz')
|
|
st_pose_euler_active = np.append(translation_active, euler_angles_active[2])
|
|
init_x = normalize_vars(st_pose_euler_active, self.og_bounds)
|
|
return init_x
|
|
|
|
def update_obj(self, obj_id, pose):
|
|
self.objects[obj_id].obj_pose = pose
|
|
|
|
def cost_function(self, opt_vars, opt_ids, constraints):
|
|
# Update obj pose
|
|
for i in range(len(opt_ids)):
|
|
obj_id = opt_ids[i]
|
|
xyz_yaw = unnormalize_vars(opt_vars[i * 4:i * 4 + 4], self.og_bounds)
|
|
|
|
pose = pose2mat([xyz_yaw[:3], euler2quat(np.append([0, 0], xyz_yaw[3]))])
|
|
self.update_obj(obj_id, pose)
|
|
|
|
# Calculate cost for each constraint
|
|
cost = 0
|
|
for info in constraints:
|
|
active_id, passive_id, constraint = info['active'], info['passive'], info['constraint']
|
|
cost += get_cost_func(constraint)(
|
|
self.objects[active_id],
|
|
self.objects[passive_id],
|
|
safe_distance=self.safe_distance
|
|
)
|
|
return cost
|
|
|
|
def __call__(self, opt_obj_ids, exist_obj_ids, constraint, visual=False):
|
|
x0 = self.init_pose(opt_obj_ids, constraint)
|
|
|
|
constraints = [constraint]
|
|
constraints.append({'active': constraint['active'], 'passive': constraint['passive'], 'constraint': 'out'})
|
|
|
|
saved_solutions = {}
|
|
# Optimization
|
|
opt_result = dual_annealing(
|
|
func=self.cost_function,
|
|
args=[opt_obj_ids, constraints],
|
|
bounds=self.norm_bounds,
|
|
maxfun=10000,
|
|
x0=x0,
|
|
no_local_search=True,
|
|
minimizer_kwargs={
|
|
'method': 'L-BFGS-B',
|
|
'options': {'maxiter': 200},
|
|
},
|
|
restart_temp_ratio=2e-5 #
|
|
)
|
|
|
|
if opt_result.success:
|
|
print("Optimization successful.")
|
|
return opt_obj_ids
|
|
else:
|
|
print("Optimization failed:", opt_result.message)
|
|
saved_solutions = None
|
|
|
|
return saved_solutions
|
|
|
|
|
|
|
|
|
|
|
|
|