| import os |
| import numpy as np |
| import torch |
| from PIL import Image, ImageDraw |
| from typing import Optional, Union |
| import random |
| import yaml |
| import argparse |
| import imageio.v2 as imageio |
| from src.utils.utils import draw_obstacles_pixel, create_mask_from_binary_csv, apply_random_obstacle_variations_fix_seed, apply_shelf_obstacle_variations_fix_seed, apply_room_obstacle_variations_fix_seed |
| from scipy import interpolate |
|
|
| def make_position_seed(seed_base, obstacle_var_idx, position_var_idx): |
| """Compute a reproducible per-run seed.""" |
| if seed_base is None: |
| seed_base = 0 |
| return int(seed_base + obstacle_var_idx * 1000 + position_var_idx) |
|
|
| def set_global_seed(seed: int): |
| """Set random seed for reproducibility""" |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
|
|
|
|
| def normalized_to_pixel(coords, img_size): |
| c = np.asarray(coords) |
| px = ((np.clip(c, -1, 1) + 1) * 0.5 * (img_size - 1)).round().astype(int) |
| return px |
|
|
|
|
| def generate_multi_goals_and_background( |
| map_name, num_goals, goal_prompts, bg_image, obstacle_tensor, |
| obstacle_ood, goal_bounds, img_size, eval_config, device |
| ): |
|
|
| if map_name == "ood": |
| background_with_obstacle = draw_obstacles_pixel(bg_image.copy(), obstacle_ood.squeeze(1)) |
|
|
| apple_indices = [i for i, prompt in enumerate(goal_prompts) if "apple" in prompt.lower()] |
| non_apple_indices = [i for i in range(num_goals) if i not in apple_indices] |
|
|
| multi_goals = torch.zeros(1, num_goals, 2, device=device, dtype=torch.float32) |
|
|
| if len(apple_indices) > 0: |
| fixed_pos = eval_config['ood']['apple_fixed_position'] |
| multi_goals[0, apple_indices[0], 0] = fixed_pos[0] |
| multi_goals[0, apple_indices[0], 1] = fixed_pos[1] |
|
|
| random_goal_indices = non_apple_indices + apple_indices[1:] |
| if len(random_goal_indices) > 0: |
| ood_goal_dist = eval_config['evaluation']['goal_generation']['min_distance_between_goals'] |
| ood_clearance = eval_config['ood']['goal_generation'].get('obstacle_clearance_pixels', |
| eval_config['evaluation']['obstacle_clearance_pixels']) |
| random_goals = gen_goals( |
| goal_bounds, (len(random_goal_indices), 1), img_size, dist=ood_goal_dist, |
| obstacles=obstacle_tensor, device=device, obstacle_clearance_pixels=ood_clearance |
| ) |
| for i, idx in enumerate(random_goal_indices): |
| multi_goals[0, idx, :] = random_goals[0, i, :] |
| else: |
| background_with_obstacle = draw_obstacles_pixel(bg_image.copy(), obstacle_tensor.squeeze(1)) |
| goal_dist = eval_config['evaluation']['goal_generation']['min_distance_between_goals'] |
| clearance = eval_config['evaluation']['obstacle_clearance_pixels'] |
| multi_goals = gen_goals( |
| goal_bounds, (num_goals, 1), img_size, dist=goal_dist, |
| obstacles=obstacle_tensor, device=device, obstacle_clearance_pixels=clearance |
| ) |
|
|
| return multi_goals, background_with_obstacle |
|
|
|
|
| def get_distance(a, b): |
| return torch.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2) |
|
|
|
|
| def is_valid_goals(pixel_x, pixel_y, obstacles, batch_idx, clearance=0): |
|
|
| if obstacles.dim() == 4: |
| obstacles = obstacles.squeeze(1) |
| |
| B, H, W = obstacles.shape |
| |
| if not (0 <= pixel_x < W and 0 <= pixel_y < H): |
| return False |
|
|
| x_start = max(0, pixel_x - clearance) |
| x_end = min(W, pixel_x + clearance + 1) |
| y_start = max(0, pixel_y - clearance) |
| y_end = min(H, pixel_y + clearance + 1) |
|
|
| area_to_check = obstacles[batch_idx, x_start:x_end, y_start:y_end] |
| |
| return torch.sum(area_to_check) == 0 |
|
|
| def gen_goals( |
| bounds, |
| n:Union[tuple, int], |
| img_size: int, |
| dist:Optional[float]=0.25, |
| obstacles:Optional[torch.Tensor]=None, |
| seed: Optional[int]=None, |
| device='cuda', |
| obstacle_clearance_pixels: int = 0 |
| ): |
| if seed is not None: |
| torch.manual_seed(seed) |
| |
| assert len(bounds) == 4, f'Unappropriate map bound: {bounds}' |
| if isinstance(n, int): |
| M = 1 |
| N = n |
| else: |
| M, N = n |
| |
| goals = [] |
| for batch_idx in range(N): |
| batch_goals = [] |
| while len(batch_goals) < M: |
| x = bounds[0] + (bounds[1] - bounds[0]) * torch.rand((1, 1), dtype=torch.float32, device=device) |
| y = bounds[2] + (bounds[3] - bounds[2]) * torch.rand((1, 1), dtype=torch.float32, device=device) |
|
|
| if obstacles is not None: |
| pixel_x = ((x + 1) * 0.5 * (img_size - 1)).long().item() |
| pixel_y = ((y + 1) * 0.5 * (img_size - 1)).long().item() |
|
|
| if not is_valid_goals(pixel_x, pixel_y, obstacles, batch_idx, clearance=obstacle_clearance_pixels): |
| continue |
|
|
| valid = True |
| if dist is not None: |
| for existing_goal in batch_goals: |
| if get_distance([x, y], [existing_goal[0, 0], existing_goal[0, 1]]) < dist: |
| valid = False |
| break |
| |
| if valid: |
| batch_goals.append(torch.cat((x, y), dim=1)) |
|
|
| goals.append(torch.cat(batch_goals, dim=0)) |
| |
| return torch.stack(goals) |
|
|
|
|
| def gen_agents( |
| bounds, |
| n:Union[tuple, int], |
| img_size: int, |
| dist:Optional[float]=None, |
| obstacles:Optional[torch.Tensor]=None, |
| seed: Optional[int]=None, |
| device='cuda', |
| obstacle_clearance_pixels: int = 0, |
| goals_to_avoid: Optional[torch.Tensor] = None, |
| min_goal_dist: Optional[float] = None |
| ): |
| if seed is not None: |
| torch.manual_seed(seed) |
| |
| assert len(bounds) == 4, f'Unappropriate map bound: {bounds}' |
| if isinstance(n, int): |
| M = 1 |
| N = n |
| else: |
| M, N = n |
| |
| goals = [] |
| for batch_idx in range(N): |
| batch_goals = [] |
| while len(batch_goals) < M: |
| x = bounds[0] + (bounds[1] - bounds[0]) * torch.rand((1, 1), dtype=torch.float32, device=device) |
| y = bounds[2] + (bounds[3] - bounds[2]) * torch.rand((1, 1), dtype=torch.float32, device=device) |
|
|
| if obstacles is not None: |
| pixel_x = ((x + 1) * 0.5 * (img_size - 1)).long().item() |
| pixel_y = ((y + 1) * 0.5 * (img_size - 1)).long().item() |
| if not is_valid_goals(pixel_x, pixel_y, obstacles, batch_idx, clearance=obstacle_clearance_pixels): |
| continue |
|
|
| valid = True |
| if dist is not None: |
| for existing_agent in batch_goals: |
| if get_distance([x, y], [existing_agent[0, 0], existing_agent[0, 1]]) < dist: |
| valid = False |
| break |
| if not valid: |
| continue |
|
|
| if goals_to_avoid is not None and min_goal_dist is not None: |
| all_goals_for_batch = goals_to_avoid[batch_idx] |
| for i in range(all_goals_for_batch.shape[0]): |
| existing_goal = all_goals_for_batch[i] |
| |
| if get_distance([x, y], [existing_goal[0], existing_goal[1]]) < min_goal_dist: |
| valid = False |
| break |
| if not valid: |
| continue |
| |
| batch_goals.append(torch.cat((x, y), dim=1)) |
|
|
| goals.append(torch.cat(batch_goals, dim=0)) |
| |
| return torch.stack(goals) |
|
|
|
|
|
|
| def overlay_goals_with_prompts_add_background(img, img_size, objs, pos, goal_prompts, goal_size): |
|
|
| assert pos.dim() == 3 |
| |
| if len(img) != pos.shape[0]: |
| img = img * pos.shape[0] |
|
|
| original_backgrounds = [im.copy() for im in img] |
| Orig_W, Orig_H = original_backgrounds[0].size |
| |
| resized_backgrounds = [im.resize((img_size, img_size), Image.LANCZOS) for im in img] |
| W, H = resized_backgrounds[0].size |
|
|
| objs_input_list = [img.copy() for img in objs] |
| |
| objs_resized = [] |
| for i in range(len(objs_input_list)): |
| new_w_res = W // goal_size |
| new_h_res = H // goal_size |
| objs_resized.append(objs_input_list[i].resize((new_w_res, new_h_res), Image.LANCZOS)) |
| |
| objs_original_sized = [] |
| for i in range(len(objs_input_list)): |
| new_w_orig = Orig_W // goal_size |
| new_h_orig = Orig_H // goal_size |
| objs_original_sized.append(objs_input_list[i].resize((new_w_orig, new_h_orig), Image.LANCZOS)) |
|
|
| pos_pix_resized = ((1 + pos) / 2 * torch.tensor([H - 1, W - 1], device=pos.device, dtype=pos.dtype)).cpu().numpy() |
| pos_pix_original = ((1 + pos) / 2 * torch.tensor([Orig_H - 1, Orig_W - 1], device=pos.device, dtype=pos.dtype)).cpu().numpy() |
| |
| |
| imgs_resized_np = [] |
| imgs_original_pil = [] |
| batch_prompts = [] |
| |
| for i in range(pos_pix_resized.shape[0]): |
| |
| bg_resized = resized_backgrounds[i].copy() |
| bg_original = original_backgrounds[i].copy() |
|
|
| ps_resized = pos_pix_resized[i] |
| ps_original = pos_pix_original[i] |
|
|
| for j in range(ps_resized.shape[0]): |
| obj_to_paste_res = objs_resized[j] |
| w_res, h_res = obj_to_paste_res.size |
| p_res = ps_resized[j] |
| p1_res, p0_res = round(p_res[1]), round(p_res[0]) |
| bg_resized.paste(obj_to_paste_res, (p1_res - w_res // 2, p0_res - h_res // 2), obj_to_paste_res) |
|
|
| obj_to_paste_orig = objs_original_sized[j] |
| w_orig, h_orig = obj_to_paste_orig.size |
| p_orig = ps_original[j] |
| p1_orig, p0_orig = round(p_orig[1]), round(p_orig[0]) |
| bg_original.paste(obj_to_paste_orig, (p1_orig - w_orig // 2, p0_orig - h_orig // 2), obj_to_paste_orig) |
|
|
| img_np = np.array(bg_resized)[..., :3] / 255 |
| imgs_resized_np.append(img_np) |
| imgs_original_pil.append(bg_original) |
| |
| batch_prompts.append([f"{goal_prompts[j]}" for j in range(len(ps_resized))]) |
|
|
| imgs_tensor = torch.tensor(imgs_resized_np, dtype=pos.dtype, device=pos.device).permute(0, 3, 1, 2) |
| |
| return imgs_tensor, batch_prompts, imgs_original_pil |
|
|
| def draw_agents_on_map( |
| agent_positions, |
| background_image, |
| dot_size=15 |
| ): |
|
|
| mapp = background_image.copy() |
| draw = ImageDraw.Draw(mapp) |
| img_width, img_height = mapp.size |
| |
| if agent_positions.dim() == 3: |
| points = agent_positions[0].cpu().numpy() |
| else: |
| points = agent_positions.cpu().numpy() |
| |
|
|
| h_coords = (points[:, 0] + 1) / 2 * (img_height - 1) |
| w_coords = (points[:, 1] + 1) / 2 * (img_width - 1) |
| |
| agent_colors = ["red", "green", "blue", "orange", "purple", "cyan", "magenta", "yellow", "brown"] |
|
|
| for i in range(points.shape[0]): |
| xc, yc = w_coords[i], h_coords[i] |
| color = agent_colors[i % len(agent_colors)] |
|
|
| bbox = [ |
| (xc - dot_size / 2, yc - dot_size / 2), |
| (xc + dot_size / 2, yc + dot_size / 2) |
| ] |
| draw.ellipse(bbox, fill=color) |
|
|
| return [mapp] |
|
|
|
|
| def load_config(path): |
| with open(path, 'r') as file: |
| config = yaml.safe_load(file) |
| return config |
|
|
| def str2bool(v): |
| if isinstance(v, bool): |
| return v |
| if v.lower() in ('yes', 'true', 't', 'y', '1'): |
| return True |
| elif v.lower() in ('no', 'false', 'f', 'n', '0'): |
| return False |
| else: |
| raise argparse.ArgumentTypeError('Boolean value expected.') |
|
|
| def collision_grad(x, d_min=0.084, eps=1e-8): |
| diff = x.unsqueeze(2) - x.unsqueeze(1) |
| dist = torch.linalg.norm(diff, dim=-1) + eps |
| mask = dist < d_min |
| grad_mag = torch.where(mask, (d_min - dist) / dist, 0.) |
| grad = (grad_mag.unsqueeze(-1) * diff).sum(dim=2) |
| return grad |
|
|
|
|
| def compute_path_length_from_pos(trajs_pos, real_world_meters=18.0): |
| METERS_PER_NORMALIZED_UNIT = real_world_meters / 2.0 |
| |
| if trajs_pos.ndim == 4: |
| path_length = torch.linalg.norm(torch.diff(trajs_pos, dim=2), dim=-1).sum(2) |
| elif trajs_pos.ndim == 3: |
| path_length = torch.linalg.norm(torch.diff(trajs_pos, dim=1), dim=-1).sum(1) |
| else: |
| raise ValueError(f"Expected trajs_pos to be 3D or 4D, got {trajs_pos.ndim}D") |
| |
| path_length_meters = path_length * METERS_PER_NORMALIZED_UNIT |
| |
| return path_length_meters |
|
|
|
|
| def resample_trajectory_constant_velocity(trajs_pos, num_samples=None): |
|
|
| if trajs_pos.ndim == 4: |
| B, N, T, D = trajs_pos.shape |
| elif trajs_pos.ndim == 3: |
| trajs_pos = trajs_pos.unsqueeze(0) |
| B, N, T, D = trajs_pos.shape |
| squeeze_output = True |
| else: |
| raise ValueError(f"Expected trajs_pos to be 3D or 4D, got {trajs_pos.ndim}D") |
| |
| if num_samples is None: |
| num_samples = T |
| |
| resampled_trajs = torch.zeros(B, N, num_samples, D, device=trajs_pos.device) |
| |
| for b in range(B): |
| for n in range(N): |
| traj = trajs_pos[b, n].cpu().numpy() |
| |
| diffs = np.diff(traj, axis=0) |
| segment_lengths = np.linalg.norm(diffs, axis=1) |
| arc_length = np.concatenate([[0], np.cumsum(segment_lengths)]) |
| |
| arc_length_diffs = np.diff(arc_length) |
| if np.any(arc_length_diffs <= 1e-10): |
| indices = np.linspace(0, T - 1, num_samples) |
| for d in range(D): |
| resampled_trajs[b, n, :, d] = torch.from_numpy( |
| np.interp(indices, np.arange(T), traj[:, d]) |
| ).to(trajs_pos.device) |
| continue |
| |
| |
| if arc_length[-1] < 1e-6: |
| resampled_trajs[b, n] = trajs_pos[b, n, 0].unsqueeze(0).repeat(num_samples, 1) |
| continue |
| |
| cs_x = interpolate.CubicSpline(arc_length, traj[:, 0]) |
| cs_y = interpolate.CubicSpline(arc_length, traj[:, 1]) |
| |
| new_arc_length = np.linspace(0, arc_length[-1], num_samples) |
| resampled_x = cs_x(new_arc_length) |
| resampled_y = cs_y(new_arc_length) |
| |
| resampled_trajs[b, n, :, 0] = torch.from_numpy(resampled_x).to(trajs_pos.device) |
| resampled_trajs[b, n, :, 1] = torch.from_numpy(resampled_y).to(trajs_pos.device) |
| |
| if 'squeeze_output' in locals() and squeeze_output: |
| resampled_trajs = resampled_trajs.squeeze(0) |
| |
| return resampled_trajs |
|
|
|
|
| def compute_energy_consumption(trajs_pos, wheel_radius=0.1, wheel_base=0.33, real_world_meters=18.0): |
| METERS_PER_NORMALIZED_UNIT = real_world_meters / 2.0 |
| k1_left = 0.0203 |
| k2_left = 0.148 |
| k3_left = 0.0014 |
| k4_left = 0.102 |
| k5_left = 0.0017 |
| P_static_left = 0.129 |
| |
| k1_right = 0.0203 |
| k2_right = 0.148 |
| k3_right = 0.0014 |
| k4_right = 0.0949 |
| k5_right = 0.0017 |
| P_static_right = 0.111 |
| |
| if trajs_pos.ndim == 4: |
| B, N, T, D = trajs_pos.shape |
| elif trajs_pos.ndim == 3: |
| trajs_pos = trajs_pos.unsqueeze(0) |
| B, N, T, D = trajs_pos.shape |
| squeeze_output = True |
| else: |
| raise ValueError(f"Expected trajs_pos to be 3D or 4D, got {trajs_pos.ndim}D") |
| |
| path_lengths_meters = compute_path_length_from_pos(trajs_pos, real_world_meters=real_world_meters) |
| |
| delta_t = 1.0 |
| total_time = T * delta_t |
| v_constant_meters = path_lengths_meters / total_time |
| |
| dx = torch.diff(trajs_pos[..., 1], dim=2) * METERS_PER_NORMALIZED_UNIT |
| dy = torch.diff(trajs_pos[..., 0], dim=2) * METERS_PER_NORMALIZED_UNIT |
| |
| ddx = torch.diff(dx, dim=2) |
| ddy = torch.diff(dy, dim=2) |
|
|
| dx_mid = (dx[:, :, :-1] + dx[:, :, 1:]) / 2 |
| dy_mid = (dy[:, :, :-1] + dy[:, :, 1:]) / 2 |
| v_const = v_constant_meters.unsqueeze(-1) |
| denom = (v_const**3).clamp(min=1e-6) |
| kappa = (dx_mid * ddy - dy_mid * ddx) / denom |
| omega = v_const.squeeze(-1).unsqueeze(-1) * kappa |
| omega = torch.nn.functional.pad(omega, (0, 1), value=0) |
| |
| omega_dot = torch.diff(omega, dim=2) / delta_t |
| omega_dot = torch.nn.functional.pad(omega_dot, (0, 1), value=0) |
| |
| energy = torch.zeros(B, N, device=trajs_pos.device) |
| |
| for t in range(T - 1): |
| v = v_constant_meters |
| omega_t = omega[:, :, t] |
| alpha_t = omega_dot[:, :, t] |
| |
| theta_dot_left = (2 * v + wheel_base * omega_t) / (2 * wheel_radius) |
| theta_dot_right = (2 * v - wheel_base * omega_t) / (2 * wheel_radius) |
| |
| theta_ddot_left = wheel_base * alpha_t / (2 * wheel_radius) |
| theta_ddot_right = -wheel_base * alpha_t / (2 * wheel_radius) |
| |
| power_left = (k1_left * theta_ddot_left**2 + |
| k2_left * theta_dot_left * theta_ddot_left + |
| k3_left * theta_ddot_left + |
| k4_left * theta_dot_left**2 + |
| k5_left * theta_dot_left + |
| P_static_left) |
| |
| power_right = (k1_right * theta_ddot_right**2 + |
| k2_right * theta_dot_right * theta_ddot_right + |
| k3_right * theta_ddot_right + |
| k4_right * theta_dot_right**2 + |
| k5_right * theta_dot_right + |
| P_static_right) |
| |
| energy += (power_left + power_right) * delta_t |
| |
| if 'squeeze_output' in locals() and squeeze_output: |
| energy = energy.squeeze(0) |
| |
| return energy |
|
|
|
|
| def compute_trajectory_smoothness(trajs_pos, real_world_meters=18.0): |
| METERS_PER_NORMALIZED_UNIT = real_world_meters / 2.0 |
|
|
| if trajs_pos.ndim == 4: |
| B, N, T, D = trajs_pos.shape |
| elif trajs_pos.ndim == 3: |
| trajs_pos = trajs_pos.unsqueeze(0) |
| B, N, T, D = trajs_pos.shape |
| squeeze_output = True |
| else: |
| raise ValueError(f"Expected trajs_pos to be 3D or 4D, got {trajs_pos.ndim}D") |
| |
|
|
| path_lengths_meters = compute_path_length_from_pos(trajs_pos, real_world_meters=real_world_meters) |
|
|
| v_const = path_lengths_meters / (T - 1) |
| dx = torch.diff(trajs_pos[..., 1], dim=2) * METERS_PER_NORMALIZED_UNIT |
| dy = torch.diff(trajs_pos[..., 0], dim=2) * METERS_PER_NORMALIZED_UNIT |
| |
| ddx = torch.diff(dx, dim=2) |
| ddy = torch.diff(dy, dim=2) |
| |
| dx_mid = (dx[:, :, :-1] + dx[:, :, 1:]) / 2 |
| dy_mid = (dy[:, :, :-1] + dy[:, :, 1:]) / 2 |
| |
| denom = (v_const.unsqueeze(-1)**3).clamp(min=1e-6) |
| kappa = (dx_mid * ddy - dy_mid * ddx) / denom |
|
|
| kappa_diff = torch.diff(kappa, dim=2) |
| |
| smoothness = torch.mean((kappa_diff**2), dim=2) |
|
|
| if 'squeeze_output' in locals() and squeeze_output: |
| smoothness = smoothness.squeeze(0) |
| |
| return smoothness |
|
|
|
|
| def compute_safety_margins(trajs_pos, obstacle_tensor, robot_radius=0.02775, img_size=96, real_world_meters=18.0): |
| METERS_PER_NORMALIZED_UNIT = real_world_meters / 2.0 |
| |
| if trajs_pos.ndim == 4: |
| B, N, T, D = trajs_pos.shape |
| elif trajs_pos.ndim == 3: |
| trajs_pos = trajs_pos.unsqueeze(0) |
| B, N, T, D = trajs_pos.shape |
| else: |
| raise ValueError(f"Expected trajs_pos to be 3D or 4D, got {trajs_pos.ndim}D") |
| |
| if obstacle_tensor.dim() == 4: |
| obstacle_tensor = obstacle_tensor.squeeze(1) |
|
|
| |
| min_dist_to_obs = float('inf') |
| for b in range(B): |
| obstacle_pixels = torch.nonzero(obstacle_tensor[b], as_tuple=False) |
| |
| if obstacle_pixels.numel() == 0: |
| continue |
|
|
| obs_y_normalized = (obstacle_pixels[:, 0].float() / (img_size - 1)) * 2 - 1 |
| obs_x_normalized = (obstacle_pixels[:, 1].float() / (img_size - 1)) * 2 - 1 |
| |
| obs_normalized = torch.stack([obs_y_normalized, obs_x_normalized], dim=1) |
| |
| for t in range(T): |
| for n in range(N): |
| robot_pos = trajs_pos[b, n, t] |
| |
| distances = torch.linalg.norm(obs_normalized - robot_pos.unsqueeze(0), dim=1) |
| |
| min_idx = distances.argmin() |
| raw_dist = distances[min_idx].item() |
| min_dist = raw_dist - robot_radius |
| |
| min_dist_meters = min_dist * METERS_PER_NORMALIZED_UNIT |
| |
| if min_dist_meters < min_dist_to_obs: |
| min_dist_to_obs = min_dist_meters |
| |
| min_dist_to_robots = float('inf') |
| |
| for t in range(T): |
| for b in range(B): |
| for i in range(N): |
| for j in range(i+1, N): |
| dist = torch.linalg.norm(trajs_pos[b, i, t] - trajs_pos[b, j, t]).item() |
| dist_meters = dist * METERS_PER_NORMALIZED_UNIT |
| if dist_meters < min_dist_to_robots: |
| min_dist_to_robots = dist_meters |
| |
| if min_dist_to_obs == float('inf'): |
| min_dist_to_obs = 1.0 * METERS_PER_NORMALIZED_UNIT |
| if min_dist_to_robots == float('inf'): |
| min_dist_to_robots = 2.0 * METERS_PER_NORMALIZED_UNIT |
| |
| return min_dist_to_obs, min_dist_to_robots |
|
|
|
|
| def get_text_embeds(texts, tokenizer, text_encoder, device): |
| """Get CLIP text embeddings""" |
| tokens = tokenizer(texts, padding="max_length", max_length=10, return_tensors="pt") |
| with torch.no_grad(): |
| outputs = text_encoder(**tokens) |
| return outputs.last_hidden_state.to(device) |
|
|
|
|
| def collision_mask(states, obstacle): |
| B, N, _ = states.shape |
| H, W = obstacle.shape[1], obstacle.shape[2] |
|
|
| mapped_states = ((states + 1) * 47.5).long() |
|
|
| y_coords = mapped_states[:, :, 0] |
| x_coords = mapped_states[:, :, 1] |
|
|
| y_coords = y_coords.clamp(0, H - 1) |
| x_coords = x_coords.clamp(0, W - 1) |
|
|
| batch_indices = torch.arange(B, device=states.device).unsqueeze(1).expand(B, N) |
|
|
| collision_mask_tensor = obstacle[batch_indices, y_coords, x_coords] |
|
|
| return collision_mask_tensor.bool() |
|
|
|
|
| def visualize_final_trajectories( |
| trajectories, |
| background_img, |
| output_path, |
| line_width=3, |
| state_dot_size=14, |
| start_dot_size=14, |
| end_dot_size=14, |
| agent_colors=None |
| ): |
|
|
| if agent_colors is None: |
| agent_colors = ["red", "green", "blue", "orange", "purple", "cyan", "magenta", "yellow", "brown"] |
| |
| mapp = background_img.copy() |
| draw = ImageDraw.Draw(mapp) |
| |
| B, N, T, _ = trajectories.shape |
| img_height, img_width = mapp.size |
|
|
| agent_trajectories = trajectories[0].cpu().numpy() |
|
|
| for i in range(N): |
| color = agent_colors[i % len(agent_colors)] |
| |
| path_points = agent_trajectories[i] |
| h_coords = (path_points[:, 0] + 1) / 2 * img_height |
| w_coords = (path_points[:, 1] + 1) / 2 * img_width |
| pixel_path = list(zip(w_coords, h_coords)) |
| |
| if len(pixel_path) > 1: |
| draw.line(pixel_path, fill=color, width=line_width) |
| |
| start_x, start_y = pixel_path[0] |
| end_x, end_y = pixel_path[-1] |
| |
| draw.ellipse( |
| [(start_x - start_dot_size/2, start_y - start_dot_size/2), |
| (start_x + start_dot_size/2, start_y + start_dot_size/2)], |
| fill="white", outline=color, width=2 |
| ) |
| |
| draw.ellipse( |
| [(end_x - end_dot_size/2, end_y - end_dot_size/2), |
| (end_x + end_dot_size/2, end_y + end_dot_size/2)], |
| fill=color, outline="black", width=2 |
| ) |
| |
| mapp.save(output_path) |
|
|
|
|
| def create_trajectory_video( |
| trajectory_history, |
| background_img, |
| output_path, |
| dot_size=10, |
| fps=50, |
| repeat_frames=3 |
| ): |
| frames = [] |
| trajectory_reversed = trajectory_history |
| |
| for x in trajectory_reversed: |
| frame_img = background_img.copy() |
| |
| for agent_img in draw_agents_on_map( |
| agent_positions=x.cpu(), |
| background_image=frame_img, |
| dot_size=dot_size |
| ): |
| frame_array = np.array(agent_img) |
| frames.append(frame_array) |
| break |
| |
| slow_frames = [] |
| for f in frames: |
| slow_frames.extend([f] * repeat_frames) |
| |
| imageio.mimsave(output_path, slow_frames, fps=fps, codec='libx264', quality=8) |
|
|
|
|
| def load_obstacle_mask_from_config(map_name, img_size, device, config, project_root): |
|
|
| mask_paths = config['mask_paths'] |
| obs_config = config['obstacle_variations'] |
| |
| if map_name == "ood": |
| csv_path = os.path.join(project_root, mask_paths['ood']) |
| base_obstacle_tensor = create_mask_from_binary_csv( |
| csv_path=csv_path, img_size=img_size, device=device |
| ) |
| return base_obstacle_tensor.unsqueeze(1) |
| |
| elif map_name == "room": |
| csv_h_path = os.path.join(project_root, mask_paths['room_horizontal']) |
| csv_v_path = os.path.join(project_root, mask_paths['room_vertical']) |
| |
| h_tensor = create_mask_from_binary_csv(csv_h_path, img_size, device) |
| v_tensor = create_mask_from_binary_csv(csv_v_path, img_size, device) |
| |
| varied_tensors = apply_room_obstacle_variations_fix_seed( |
| horizontal_mask_tensor=h_tensor, |
| vertical_mask_tensor=v_tensor, |
| num_variations=obs_config['num_variations'], |
| device=device, |
| seed=obs_config['seed'], |
| scale_range=tuple(obs_config['scale_range']), |
| offset_range_pixels=obs_config['offset_range_pixels'] |
| ) |
| |
| elif map_name == "shelf": |
| csv_path = os.path.join(project_root, mask_paths['shelf']) |
| base_obstacle_tensor = create_mask_from_binary_csv( |
| csv_path=csv_path, img_size=img_size, device=device |
| ) |
| cropped_obstacle_tensor = base_obstacle_tensor.clone() |
|
|
| shelf_config = obs_config['shelf'] |
| for crop_region in shelf_config['crop_regions']: |
| y_start, y_end, x_start, x_end = crop_region |
| cropped_obstacle_tensor[0, y_start:y_end, x_start:x_end] = 0 |
| |
| varied_tensors = apply_shelf_obstacle_variations_fix_seed( |
| base_obstacle_mask_tensor=cropped_obstacle_tensor, |
| num_variations=obs_config['num_variations'], |
| device=device, |
| seed=obs_config['seed'], |
| scale_range=tuple(shelf_config['scale_range']), |
| offset_range_pixels=obs_config['offset_range_pixels'], |
| min_obstacle_area_ratio=obs_config['min_obstacle_area_ratio'], |
| num_special_variations=shelf_config['num_special_variations'], |
| special_vertical_offset=shelf_config['special_vertical_offset'] |
| ) |
| else: |
| csv_path = os.path.join(project_root, mask_paths[map_name]) |
| base_tensor = create_mask_from_binary_csv(csv_path, img_size, device) |
| |
| varied_tensors = apply_random_obstacle_variations_fix_seed( |
| base_obstacle_mask_tensor=base_tensor, |
| num_variations=obs_config['num_variations'], |
| device=device, |
| seed=obs_config['seed'], |
| scale_range=tuple(obs_config['scale_range']), |
| offset_range_pixels=obs_config['offset_range_pixels'], |
| min_obstacle_area_ratio=obs_config['min_obstacle_area_ratio'] |
| ) |
| |
| return varied_tensors |
|
|
|
|
| def create_model_from_checkpoint(img_size, device, checkpoint_path, Unet): |
| |
| class Unet2D(Unet): |
| def __init__(self, dim, out_dim, dim_mults=(1, 2, 4, 8)): |
| super().__init__(dim=dim, out_dim=out_dim, dim_mults=dim_mults) |
|
|
| def forward(self, obs, t, prompt): |
| predicted_fields = super().forward(obs, t, prompt=prompt) |
| return predicted_fields |
|
|
| model = Unet2D( |
| dim=img_size, |
| out_dim=2, |
| dim_mults=(1, 2, 4, 8), |
| ).to(device) |
| |
| |
| checkpoint = torch.load(checkpoint_path, map_location=device) |
| if isinstance(checkpoint, dict) and 'model' in checkpoint: |
| model.load_state_dict(checkpoint['model']) |
| else: |
| model.load_state_dict(checkpoint) |
| |
| model.eval() |
| return model |
|
|
|
|
| def load_evaluation_config(config_path): |
|
|
| with open(config_path, 'r') as f: |
| config = yaml.safe_load(f) |
| return config |
|
|
|
|
|
|
| def inference_combined( |
| model, diffusion, obs1, prompt_list, x_T, obstacle_tensor, |
| frozen, img_size, device, bilinear_interpolate_samples, |
| config, seed: int = None, use_gaussian: bool = False, step_size: float = 0.025, |
| use_inter_robot_collision: bool = True |
| ): |
| if seed is not None: |
| torch.manual_seed(seed) |
|
|
| inf_config = config['inference'] |
| epsilon = inf_config['epsilon'] |
| k1 = inf_config['k1'] |
| k2 = inf_config['k2'] |
| beta0 = inf_config['beta0'] |
| beta_pow = inf_config['beta_pow'] |
| tau_coll = inf_config['tau_coll'] |
| anneal_it = inf_config['anneal_iterations'] |
|
|
| robot_diameter = config['evaluation']['robot_diameter'] |
|
|
| T = diffusion.noise_steps |
| x = x_T.clone() |
| B, N, D = x.shape |
|
|
| prompts_N = torch.cat(prompt_list, dim=0) |
| prompts_BN = prompts_N.unsqueeze(0).repeat(B, 1, 1, 1) |
| prompts_flat = prompts_BN.view(-1, 10, 512) |
|
|
| obs_flat = obs1.repeat_interleave(N, 0) |
|
|
| use_precompute = (N == 3) |
|
|
| if use_precompute: |
| all_t_steps = torch.arange(1, T + 1, device=device) |
| t_full_batch = all_t_steps.unsqueeze(1).repeat(1, B*N).view(-1) |
|
|
| obs_full_batch = obs_flat.unsqueeze(0).repeat(T, 1, 1, 1, 1).view(-1, *obs_flat.shape[1:]) |
| prompts_full_batch = prompts_flat.unsqueeze(0).repeat(T, 1, 1, 1).view(-1, *prompts_flat.shape[1:]) |
|
|
| all_predicted_fields_flat = model(obs_full_batch, t_full_batch, prompts_full_batch) |
| field_shape = all_predicted_fields_flat.shape[1:] |
| all_predicted_fields = all_predicted_fields_flat.view(T, B*N, *field_shape) |
| |
| trajectory_history = [] |
|
|
| for i in reversed(range(1, T+1)): |
| alpha = epsilon * (diffusion.std[i-1] / diffusion.std[-1]) |
| collision_masks = collision_mask(x, obstacle_tensor).to(device) |
| movable = ~frozen & ~collision_masks |
|
|
| active_mask = (movable).view(-1) |
| idx_active = active_mask.nonzero(as_tuple=False).squeeze(-1) |
| |
| if use_precompute: |
| fields_for_t_i_all_agents = all_predicted_fields[i-1] |
| predicted_fields = fields_for_t_i_all_agents[idx_active] |
| else: |
| t = (torch.ones(1) * i).long().to(device) |
| t_flat = t.repeat_interleave(N) |
|
|
| obs_batch = obs_flat[idx_active] |
| t_batch = t_flat[idx_active] |
| prompts_batch = prompts_flat[idx_active] |
| predicted_fields = model(obs_batch, t_batch, prompts_batch) |
|
|
| for j in range(anneal_it): |
| x_prev = x.clone() |
|
|
| x_flat_all = x.view(-1, 1, D) |
| x_batch = x_flat_all[idx_active] |
|
|
| sampled_preds = bilinear_interpolate_samples(predicted_fields, x_batch) |
| score_batch = sampled_preds[..., :2] |
| score_bn = score_batch.view(-1, 1, D) |
| scores_flat = torch.zeros_like(x_flat_all) |
| scores_flat[idx_active, 0, :] = score_bn.view(-1, D) |
| scores = scores_flat.view(B, N, D) |
|
|
| grad_cost = collision_grad(x_prev, d_min=tau_coll) |
|
|
| guide_w = beta0 * (i / T) ** beta_pow |
| |
| if use_inter_robot_collision: |
| scores_hat = scores + guide_w * grad_cost |
| else: |
| scores_hat = scores |
|
|
| if not use_gaussian: |
| drift = (img_size / (2*diffusion.std[i-1])) * scores_hat * alpha**k1 |
| else: |
| drift = scores_hat * step_size |
|
|
| noise = torch.randn_like(x) * alpha**((k1+k2)/2) |
|
|
| diff_mat = x.unsqueeze(2) - x.unsqueeze(1) |
| dist_mat = torch.linalg.norm(diff_mat, dim=-1) |
| collided = (dist_mat < robot_diameter).triu(diagonal=1) |
| for b in range(B): |
| pairs = torch.nonzero(collided[b], as_tuple=False) |
| if pairs.numel() > 0: |
| agents_involved = torch.unique(pairs) |
| for n_idx in agents_involved.tolist(): |
| frozen[b, n_idx] = True |
| movable[b, n_idx] = False |
|
|
| x[movable] = x_prev[movable] + drift[movable] + noise[movable] |
| if j == 0 or j == (anneal_it // 3) or j == (anneal_it * 2 // 3) or j == (anneal_it - 1): |
| trajectory_history.append(x.clone().clamp(-0.99, 0.99)) |
|
|
| return x, ~movable, trajectory_history |
|
|
