finish task_gen

task_gen_dependencies/solver_3d.py

@@ -0,0 +1,166 @@
+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
+
+
+
+
+
+