| from __future__ import annotations |
|
|
| import math |
| from typing import Any |
|
|
| import torch |
|
|
|
|
| _ORIGINAL_GENERATE_MAPS: Any | None = None |
| _MOTIF_SETTINGS: dict[str, float | int] = {} |
|
|
|
|
| def configure_motif_mapgen(*, replace: bool = False, **kwargs: float | int) -> None: |
| if replace: |
| _MOTIF_SETTINGS.clear() |
| _MOTIF_SETTINGS.update(kwargs) |
|
|
|
|
| def _setting(cfg: Any, name: str, default: float | int) -> float | int: |
| if name in _MOTIF_SETTINGS: |
| return _MOTIF_SETTINGS[name] |
| return getattr(cfg, name, default) |
|
|
|
|
| def install_motif_mapgen() -> None: |
| """Patch the fast Torch env to use compact rhythm/motif maps. |
| |
| This is intentionally isolated under experiments/ so the base runner stays |
| usable. The patch affects only the current Python process. |
| """ |
|
|
| global _ORIGINAL_GENERATE_MAPS |
| from env import FastFloatingBeaconEnv |
|
|
| if _ORIGINAL_GENERATE_MAPS is None: |
| _ORIGINAL_GENERATE_MAPS = FastFloatingBeaconEnv.generate_maps |
| FastFloatingBeaconEnv.generate_maps = generate_motif_maps |
|
|
|
|
| def restore_original_mapgen() -> None: |
| global _ORIGINAL_GENERATE_MAPS |
| if _ORIGINAL_GENERATE_MAPS is None: |
| return |
| from env import FastFloatingBeaconEnv |
|
|
| FastFloatingBeaconEnv.generate_maps = _ORIGINAL_GENERATE_MAPS |
|
|
|
|
| def generate_motif_maps(self: Any) -> None: |
| route_len, p, n = self.route_len, self.p, self.n |
| self.platforms.zero_() |
| self.sizes.fill_(0.0) |
| base_size = torch.tensor([self.platform_x * 1.05, self.platform_y * 1.06], device=self.device) |
| min_scale = float(_setting(self.cfg, "map_platform_min_scale", 0.58)) |
| max_scale = float(_setting(self.cfg, "map_platform_max_scale", 1.36)) |
| tiny_fraction = float(_setting(self.cfg, "map_platform_tiny_fraction", 0.18)) |
| large_fraction = float(_setting(self.cfg, "map_platform_large_fraction", 0.14)) |
| aspect_jitter = float(_setting(self.cfg, "map_platform_aspect_jitter", 0.46)) |
| scale = min_scale + (max_scale - min_scale) * self.rand(n, p) |
| tiny = self.rand(n, p) < tiny_fraction |
| large = self.rand(n, p) < large_fraction |
| scale = torch.where(tiny, 0.46 + 0.18 * self.rand(n, p), scale) |
| scale = torch.where(large, 1.22 + 0.24 * self.rand(n, p), scale) |
| aspect = torch.exp((self.rand(n, p) - 0.5) * aspect_jitter).clamp(0.72, 1.38) |
| self.sizes[:, :, 0] = (base_size[0] * scale * aspect).clamp(0.44, 1.70) |
| self.sizes[:, :, 1] = (base_size[1] * scale / aspect).clamp(0.38, 1.48) |
| |
| self.sizes[:, 0] = torch.maximum(self.sizes[:, 0], base_size * 1.18) |
| self.sizes[:, route_len - 1] = torch.maximum(self.sizes[:, route_len - 1], base_size * 1.08) |
| route_tiny_fraction = float(_setting(self.cfg, "map_route_tiny_fraction", 0.0)) |
| if route_len > 3 and route_tiny_fraction > 0.0: |
| route_nodes = torch.arange(route_len, device=self.device)[None, :] |
| eligible = (route_nodes > 0) & (route_nodes < route_len - 1) |
| rhythm = ((route_nodes * 2 + self.env_i[:, None]) % 7) == 3 |
| route_tiny = eligible & (rhythm | (self.rand(n, route_len) < route_tiny_fraction * 0.20)) |
| tiny_min = max(0.34, float(_setting(self.cfg, "map_route_tiny_min_scale", 0.44))) |
| tiny_max = max(tiny_min + 0.04, float(_setting(self.cfg, "map_route_tiny_max_scale", 0.68))) |
| tiny_scale = tiny_min + (tiny_max - tiny_min) * self.rand(n, route_len) |
| tiny_aspect = torch.exp((self.rand(n, route_len) - 0.5) * 0.34).clamp(0.78, 1.28) |
| tiny_x = (base_size[0] * tiny_scale * tiny_aspect).clamp(0.38, 0.84) |
| tiny_y = (base_size[1] * tiny_scale / tiny_aspect).clamp(0.34, 0.76) |
| self.sizes[:, :route_len, 0] = torch.where(route_tiny, tiny_x, self.sizes[:, :route_len, 0]) |
| self.sizes[:, :route_len, 1] = torch.where(route_tiny, tiny_y, self.sizes[:, :route_len, 1]) |
|
|
| height_offset = max(0.0, float(getattr(self.cfg, "map_height_offset", 2.85))) |
| vertical = max(0.35, float(getattr(self.cfg, "map_vertical_scale", 1.6))) |
| max_h = max(height_offset + 0.7, float(getattr(self.cfg, "map_max_height", 7.2))) |
| edge_gap = max(0.44, float(getattr(self.cfg, "map_min_edge_gap", 0.62))) |
| edge_gap_min = max(edge_gap, float(_setting(self.cfg, "map_edge_gap_min", 0.86))) |
| edge_gap_max = max(edge_gap_min + 0.08, float(_setting(self.cfg, "map_edge_gap_max", 1.42))) |
| route_gap_scale = max(0.80, float(_setting(self.cfg, "map_route_gap_scale", 1.10))) |
| step_len_base = max(1.56, self.platform_x + edge_gap_min) |
| turn_scale = max(0.0, float(_setting(self.cfg, "map_scenic_turn_scale", 1.0))) |
| vertical_step = max(0.12, float(_setting(self.cfg, "map_scenic_vertical_step", 0.44))) * vertical |
| island_radius = max(1.18, float(_setting(self.cfg, "map_scenic_island_radius", 1.55))) |
| cleanup_iters = int(getattr(self.cfg, "map_overlap_cleanup_iters", 16)) |
| momentum_fraction = max(0.0, float(_setting(self.cfg, "map_momentum_jump_fraction", 0.0))) |
| momentum_gap_bias = min(max(float(_setting(self.cfg, "map_momentum_gap_bias", 0.82)), 0.55), 0.96) |
| redirect_period = max(0, int(_setting(self.cfg, "map_redirect_period", 0))) |
| redirect_angle = max(0.0, float(_setting(self.cfg, "map_redirect_angle", 1.05))) |
| goal_vertical_bias = max(0.0, float(_setting(self.cfg, "map_goal_vertical_bias", 0.0))) |
| goal_vertical_bias_min = min(goal_vertical_bias, max(0.0, float(_setting(self.cfg, "map_goal_vertical_bias_min", 0.0)))) |
| hairpin_period = max(0, int(_setting(self.cfg, "map_hairpin_period", 0))) |
| hairpin_angle = max(0.0, float(_setting(self.cfg, "map_hairpin_angle", 0.0))) |
| hairpin_depth = min(max(float(_setting(self.cfg, "map_hairpin_depth", 0.0)), 0.0), 1.5) |
| unintuitive_depth = min(max(float(_setting(self.cfg, "map_unintuitive_depth", 0.0)), 0.0), 1.5) |
| valley_period = max(0, int(_setting(self.cfg, "map_valley_period", 0))) |
| valley_depth = min(max(float(_setting(self.cfg, "map_valley_depth", unintuitive_depth)), 0.0), 1.5) |
| route_detour_period = max(0, int(_setting(self.cfg, "map_route_detour_period", 0))) |
| route_detour_depth = min(max(float(_setting(self.cfg, "map_route_detour_depth", 0.0)), 0.0), 1.5) |
| backtrack_period = max(0, int(_setting(self.cfg, "map_backtrack_period", 0))) |
| backtrack_depth = min(max(float(_setting(self.cfg, "map_backtrack_depth", 0.0)), 0.0), 1.5) |
| climb_direction = torch.where(self.rand(n) < 0.5, 1.0, -1.0) |
| climb_total = climb_direction * (goal_vertical_bias_min + (goal_vertical_bias - goal_vertical_bias_min) * self.rand(n)) |
| climb_per_step = climb_total / max(1, route_len - 1) |
|
|
| route_xy = torch.zeros((n, route_len, 2), device=self.device) |
| route_z = torch.zeros((n, route_len), device=self.device) |
| xy = torch.zeros((n, 2), device=self.device) |
| heading = self.rand(n) * math.tau |
| phase = self.rand(n) * math.tau |
| motif_seed = torch.randint(0, 7, (n,), device=self.device, generator=self.gen) |
| motif_offset = self.env_i % 3 |
| side_seed = torch.where((self.env_i % 2) == 0, 1.0, -1.0) |
| z = torch.full((n,), height_offset + 0.28, device=self.device) + self.rand(n) * 0.30 |
| z = (z + (-climb_total).clamp_min(0.0)).clamp(height_offset + 0.18, max_h - 0.20) |
| route_z[:, 0] = z |
| self.platforms[:, 0, 2] = z |
|
|
| turn_tables = ( |
| (0.14, 0.20, 0.14, 0.04), |
| (0.62, 0.08, -0.44, -0.18), |
| (0.04, -0.10, 0.12, -0.04), |
| (0.86, 0.24, -0.62, -0.30), |
| (0.38, -0.46, -0.28, 0.42), |
| (0.18, 0.04, -0.12, 0.08), |
| (0.50, -0.16, -0.38, 0.24), |
| ) |
| step_tables = ( |
| (1.02, 1.02, 0.96, 1.00), |
| (0.88, 1.02, 0.88, 1.04), |
| (0.92, 0.88, 0.92, 1.00), |
| (0.72, 0.86, 0.76, 0.98), |
| (0.92, 0.90, 0.94, 1.02), |
| (0.84, 0.82, 0.90, 0.98), |
| (0.90, 0.86, 0.92, 1.04), |
| ) |
| dz_tables = ( |
| (0.42, 0.35, -0.28, -0.24), |
| (0.18, 0.12, 0.22, -0.18), |
| (0.46, 0.38, 0.20, -0.24), |
| (0.34, -0.18, 0.44, -0.20), |
| (0.24, 0.34, -0.26, -0.12), |
| (0.52, 0.38, 0.18, -0.16), |
| (0.16, 0.36, -0.10, 0.12), |
| ) |
|
|
| for layer in range(1, route_len): |
| slot = (layer - 1) % 4 |
| block = (layer - 1) // 4 |
| motif = (motif_seed + int(block) + motif_offset) % 7 |
| kind = [motif == i for i in range(7)] |
| side = torch.where(((self.env_i + layer) % 2) == 0, 1.0, -1.0) * side_seed |
| turn_delta = torch.sin(phase + float(layer) * 0.73) * 0.04 |
| step_mult = torch.ones(n, device=self.device) |
| dz = torch.sin(phase + float(layer) * 1.13) * vertical_step * 0.10 |
|
|
| for i in range(7): |
| active = kind[i].float() |
| turn_delta = turn_delta + active * side * float(turn_tables[i][slot]) |
| step_mult = step_mult + active * (float(step_tables[i][slot]) - 1.0) |
| dz = dz + active * float(dz_tables[i][slot]) * vertical_step |
| dz = dz + climb_per_step |
| if redirect_period > 0 and redirect_angle > 0.0: |
| redirect_slot = (layer - 2) % redirect_period |
| if redirect_slot == 0: |
| turn_delta = turn_delta + side * redirect_angle |
| step_mult = step_mult * 0.88 |
| elif redirect_slot == 1: |
| turn_delta = turn_delta - side * redirect_angle * 0.62 |
| elif redirect_slot == 2: |
| turn_delta = turn_delta - side * redirect_angle * 0.30 |
|
|
| if hairpin_period > 0 and hairpin_angle > 0.0 and hairpin_depth > 0.0: |
| hairpin_slot = (layer - 2) % hairpin_period |
| fold_side = torch.where(((self.env_i // max(1, hairpin_period) + layer) % 2) == 0, 1.0, -1.0) * side_seed |
| if hairpin_slot == 0: |
| turn_delta = turn_delta + fold_side * hairpin_angle |
| step_mult = step_mult * (0.84 - 0.08 * hairpin_depth) |
| dz = dz + vertical_step * 0.14 * hairpin_depth |
| elif hairpin_slot == 1: |
| turn_delta = turn_delta + fold_side * hairpin_angle * 0.42 |
| step_mult = step_mult * (0.76 - 0.06 * hairpin_depth) |
| dz = dz - vertical_step * 0.10 * hairpin_depth |
| elif hairpin_slot == 2: |
| turn_delta = turn_delta - fold_side * hairpin_angle * 0.78 |
| step_mult = step_mult * (0.86 - 0.05 * hairpin_depth) |
| dz = dz + vertical_step * 0.18 * hairpin_depth |
| elif hairpin_slot == 3: |
| turn_delta = turn_delta - fold_side * hairpin_angle * 0.58 |
| step_mult = step_mult * 0.92 |
| dz = dz - vertical_step * 0.12 * hairpin_depth |
|
|
| if valley_period > 0 and valley_depth > 0.0: |
| valley_slot = (layer - 2) % valley_period |
| detour_side = torch.where(((self.env_i + layer // max(1, valley_period)) % 2) == 0, 1.0, -1.0) * side_seed |
| if valley_slot == 0: |
| |
| |
| turn_delta = turn_delta + detour_side * (0.42 + 0.22 * unintuitive_depth) |
| step_mult = step_mult * (0.88 - 0.04 * valley_depth) |
| dz = dz - vertical_step * (0.42 + 0.18 * valley_depth) |
| elif valley_slot == 1: |
| turn_delta = turn_delta + detour_side * (0.24 + 0.14 * unintuitive_depth) |
| step_mult = step_mult * 0.84 |
| dz = dz - vertical_step * (0.26 + 0.16 * valley_depth) |
| elif valley_slot == 2: |
| turn_delta = turn_delta - detour_side * 0.18 |
| step_mult = step_mult * 0.78 |
| dz = dz + vertical_step * (0.10 + 0.14 * valley_depth) |
| elif valley_slot == 3: |
| turn_delta = turn_delta - detour_side * (0.38 + 0.18 * unintuitive_depth) |
| step_mult = step_mult * 0.88 |
| dz = dz + vertical_step * (0.38 + 0.18 * valley_depth) |
| elif valley_slot == 4: |
| turn_delta = turn_delta - detour_side * (0.18 + 0.12 * unintuitive_depth) |
| step_mult = step_mult * 0.94 |
| dz = dz + vertical_step * (0.28 + 0.14 * valley_depth) |
|
|
| if route_detour_period > 0 and route_detour_depth > 0.0: |
| detour_slot = (layer - 2) % route_detour_period |
| detour_side = torch.where(((self.env_i // max(1, route_detour_period) + layer) % 2) == 0, 1.0, -1.0) * side_seed |
| if detour_slot == 0: |
| |
| |
| turn_delta = turn_delta + detour_side * (0.72 + 0.30 * route_detour_depth) |
| step_mult = step_mult * (0.86 - 0.04 * route_detour_depth) |
| dz = dz - vertical_step * (0.22 + 0.14 * route_detour_depth) |
| elif detour_slot == 1: |
| turn_delta = turn_delta + detour_side * (0.34 + 0.16 * route_detour_depth) |
| step_mult = step_mult * 0.86 |
| dz = dz + vertical_step * (0.08 + 0.08 * route_detour_depth) |
| elif detour_slot == 2: |
| turn_delta = turn_delta - detour_side * (0.64 + 0.22 * route_detour_depth) |
| step_mult = step_mult * 0.90 |
| dz = dz + vertical_step * (0.24 + 0.12 * route_detour_depth) |
|
|
| if backtrack_period > 0 and backtrack_depth > 0.0: |
| back_slot = (layer - 2) % backtrack_period |
| back_side = torch.where(((self.env_i // max(1, backtrack_period) + layer) % 2) == 0, 1.0, -1.0) * side_seed |
| if back_slot == 0: |
| |
| |
| turn_delta = turn_delta + back_side * (2.02 + 0.28 * backtrack_depth) |
| step_mult = step_mult * (0.88 + 0.06 * backtrack_depth) |
| dz = dz - vertical_step * (0.16 + 0.10 * backtrack_depth) |
| elif back_slot == 1: |
| turn_delta = turn_delta + back_side * (0.68 + 0.16 * backtrack_depth) |
| step_mult = step_mult * (0.92 + 0.04 * backtrack_depth) |
| dz = dz + vertical_step * (0.10 + 0.10 * backtrack_depth) |
| elif back_slot == 2: |
| turn_delta = turn_delta - back_side * (1.54 + 0.20 * backtrack_depth) |
| step_mult = step_mult * (0.94 + 0.03 * backtrack_depth) |
| dz = dz + vertical_step * (0.26 + 0.12 * backtrack_depth) |
| elif back_slot == 3: |
| turn_delta = turn_delta - back_side * (0.54 + 0.18 * backtrack_depth) |
| step_mult = step_mult * 0.88 |
| dz = dz - vertical_step * 0.08 * backtrack_depth |
|
|
| heading = heading + turn_delta * turn_scale |
| direction = torch.stack((torch.cos(heading), torch.sin(heading)), dim=1) |
| support = ( |
| (direction.abs() * self.sizes[:, layer - 1]).sum(dim=1) * 0.5 |
| + (direction.abs() * self.sizes[:, layer]).sum(dim=1) * 0.5 |
| ) |
| target_edge_gap = ( |
| edge_gap_min + (edge_gap_max - edge_gap_min) * self.rand(n) |
| ) * step_mult.clamp(0.72, 1.18) * route_gap_scale |
| climb_trade = (climb_per_step.clamp_min(0.0) / 0.42).clamp(0.0, 0.42) |
| target_edge_gap = target_edge_gap * (1.0 - climb_trade) |
| desired_center = support + target_edge_gap |
| landing_pad = self.sizes[:, layer].amin(dim=1) * 0.34 + self.sizes[:, layer - 1].amin(dim=1) * 0.16 + float(self.radius) * 0.75 |
| horizontal_speed = float(self.max_speed) * float(self.sprint_speed_mult) * float(getattr(self.cfg, "map_reach_scale", 0.82)) |
| needed_flight = ((desired_center - landing_pad + 0.44) / max(horizontal_speed, 1e-6)).clamp_min(0.05) |
| max_dz_for_desired_gap = float(self.jump_velocity) * needed_flight - 0.5 * float(self.gravity) * needed_flight.square() |
| dz = (dz + (self.rand(n) - 0.5) * vertical_step * 0.20).clamp(-0.72, 0.62) |
| dz = torch.minimum(dz, max_dz_for_desired_gap - 0.06) |
| next_z = (z + dz).clamp(height_offset + 0.18, max_h) |
| dz_actual = next_z - z |
| jump_v = float(self.jump_velocity) |
| gravity = float(self.gravity) |
| disc = jump_v * jump_v - 2.0 * gravity * dz_actual |
| flight = (jump_v + torch.sqrt(disc.clamp_min(0.0))) / gravity |
| max_center = float(self.max_speed) * float(self.sprint_speed_mult) * flight * float(getattr(self.cfg, "map_reach_scale", 0.82)) + landing_pad - 0.42 |
| min_center = support + edge_gap_min * 0.82 |
| safe_min = torch.minimum(min_center, max_center).clamp_min(support + edge_gap_min * 0.35) |
| center_dist = torch.minimum(desired_center, max_center).clamp_min(safe_min) |
| if momentum_fraction > 0.0: |
| rhythm = ((self.env_i + layer * 3) % 8) == 2 |
| momentum = rhythm | (self.rand(n) < momentum_fraction * 0.18) |
| hard_center = safe_min + (max_center - safe_min).clamp_min(0.0) * momentum_gap_bias |
| valid_hard = max_center > safe_min + 0.35 |
| center_dist = torch.where(momentum & valid_hard, torch.maximum(center_dist, hard_center), center_dist) |
| step = direction * center_dist[:, None] |
| xy = xy + step |
| z = next_z |
|
|
| route_xy[:, layer] = xy |
| route_z[:, layer] = z |
| self.platforms[:, layer, :2] = xy |
| self.platforms[:, layer, 2] = z |
| if layer >= 2 and float(_setting(self.cfg, "map_repair_local_lookback", 0.0)) > 0.0: |
| self.repair_route_spacing(layer, min_center=max(1.38, step_len_base * 0.72)) |
| route_xy[:, layer] = self.platforms[:, layer, :2] |
|
|
| self.platforms[:, :route_len, :2] = route_xy |
| self.platforms[:, :route_len, 2] = route_z |
| goal_xy_min = max(0.0, float(_setting(self.cfg, "map_goal_xy_min", 0.0))) |
| goal_xy_per_jump_min = max(0.0, float(_setting(self.cfg, "map_goal_xy_per_jump_min", 0.0))) |
| goal_xy_min = max(goal_xy_min, goal_xy_per_jump_min * float(max(1, route_len - 1))) |
| if goal_xy_min > 0.0 and route_len > 2: |
| start_xy = route_xy[:, 0] |
| goal_vec = route_xy[:, -1] - start_xy |
| goal_dist = torch.linalg.norm(goal_vec, dim=1) |
| first_vec = route_xy[:, 1] - start_xy |
| first_dir = first_vec / torch.linalg.norm(first_vec, dim=1, keepdim=True).clamp_min(1e-4) |
| goal_dir = torch.where(goal_dist[:, None] > 1e-4, goal_vec / goal_dist[:, None].clamp_min(1e-4), first_dir) |
| stretch = (goal_xy_min - goal_dist).clamp_min(0.0) |
| ramp = torch.linspace(0.0, 1.0, route_len, device=self.device)[None, :, None] |
| route_xy = route_xy + goal_dir[:, None, :] * stretch[:, None, None] * ramp |
| self.platforms[:, :route_len, :2] = route_xy |
|
|
| route_delta_xy = route_xy[:, 1:, :2] - route_xy[:, :-1, :2] |
| route_turn_strength = torch.zeros((n, route_len), device=self.device) |
| if route_len > 3: |
| prev_vec = route_delta_xy[:, :-1] |
| next_vec = route_delta_xy[:, 1:] |
| prev_len = torch.linalg.norm(prev_vec, dim=2).clamp_min(1e-4) |
| next_len = torch.linalg.norm(next_vec, dim=2).clamp_min(1e-4) |
| prev_n = prev_vec / prev_len[:, :, None] |
| next_n = next_vec / next_len[:, :, None] |
| cross = prev_n[..., 0] * next_n[..., 1] - prev_n[..., 1] * next_n[..., 0] |
| dot = (prev_n * next_n).sum(dim=2).clamp(-1.0, 1.0) |
| corner = torch.atan2(cross.abs(), dot) |
| route_turn_strength[:, 1:-1] = corner |
| route_turn_strength = route_turn_strength.clamp_min(0.0) |
|
|
| extra_reach_exclude: torch.Tensor | None = None |
| if p > route_len: |
| extras = p - route_len |
| unlock_lever_mask: torch.Tensor | None = None |
| unlock_target_i: torch.Tensor | None = None |
| unlock_gate_count = 0 |
| inner_layers = max(1, route_len - 2) |
| offset_seed = torch.randint(0, inner_layers, (n, 1), device=self.device, generator=self.gen) |
| extra_i = torch.arange(extras, device=self.device)[None, :] |
| even_extra = (extra_i % 2) == 0 |
| odd_extra = ~even_extra |
| anchors = 1 + ((extra_i * 2 + offset_seed) % inner_layers) |
| unlock_enabled = bool(getattr(self.cfg, "unlock_platforms", False)) and route_len > 5 |
| if unlock_enabled: |
| gate_count = min(max(1, int(getattr(self.cfg, "unlock_gate_count", 1))), extras) |
| unlock_gate_count = gate_count |
| lock_span = min(max(1, int(getattr(self.cfg, "unlock_lock_span", 2))), max(1, route_len - 4)) |
| anchor_span = max(1, route_len - lock_span - 3) |
| gate_ids = torch.arange(gate_count, device=self.device)[None, :] |
| gate_anchors = 2 + torch.remainder(self.env_i[:, None] * (5 + 2 * gate_ids) + gate_ids * 3, anchor_span) |
| anchors[:, :gate_count] = gate_anchors |
| unlock_lever_mask = torch.zeros((n, extras), dtype=torch.bool, device=self.device) |
| unlock_lever_mask[:, :gate_count] = True |
| unlock_target_i = (anchors + 1).clamp_max(route_len - 2) |
| prev_i = (anchors - 1).clamp_min(0) |
| next_i = (anchors + 1).clamp_max(route_len - 1) |
| far_i = (anchors + 2).clamp_max(route_len - 1) |
| braid_span = max(2, int(_setting(self.cfg, "map_braid_bridge_span", 3))) |
| braid_target_i = (anchors + braid_span).clamp_max(route_len - 1) |
|
|
| def gather_xy(index: torch.Tensor) -> torch.Tensor: |
| return route_xy.gather(1, index[:, :, None].expand(-1, -1, 2)) |
|
|
| anchor_xy = gather_xy(anchors) |
| prev_xy = gather_xy(prev_i) |
| next_xy = gather_xy(next_i) |
| far_xy = gather_xy(far_i) |
| braid_target_xy = gather_xy(braid_target_i) |
| tangent = next_xy - prev_xy |
| tangent = tangent / torch.linalg.norm(tangent, dim=2, keepdim=True).clamp_min(1e-4) |
| perp = torch.stack((-tangent[..., 1], tangent[..., 0]), dim=2) |
| side = torch.where(((extra_i + self.env_i[:, None]) % 2) == 0, 1.0, -1.0) |
| ring = island_radius * (1.12 + 0.48 * self.rand(n, extras)) |
| ring = torch.where((extra_i % 3) == 0, ring + 0.28, ring) |
| forward = (((extra_i % 3).float() - 1.0) * 0.28 * step_len_base)[:, :, None] |
| island_xy = anchor_xy + perp * side[:, :, None] * ring[:, :, None] + tangent * forward |
| bridge_xy = 0.50 * anchor_xy + 0.50 * far_xy + perp * side[:, :, None] * (1.06 + 0.26 * self.rand(n, extras))[:, :, None] |
| terrace_xy = anchor_xy + tangent * (0.66 + 0.22 * self.rand(n, extras))[:, :, None] + perp * side[:, :, None] * (1.22 + 0.34 * self.rand(n, extras))[:, :, None] |
| bridge_mask = ((extra_i + offset_seed) % 4) == 0 |
| terrace_mask = ((extra_i + offset_seed) % 5) == 0 |
| direct_decoy_fraction = max(0.0, float(_setting(self.cfg, "map_direct_decoy_fraction", 0.0))) |
| direct_decoy = self.rand(n, extras) < direct_decoy_fraction |
| cluster_extra_fraction = max(0.0, float(_setting(self.cfg, "map_cluster_extra_fraction", 0.0))) |
| cluster_extra = self.rand(n, extras) < cluster_extra_fraction |
| decoy_corridor_fraction = max(0.0, float(_setting(self.cfg, "map_decoy_corridor_fraction", 0.0))) |
| decoy_corridor = self.rand(n, extras) < decoy_corridor_fraction |
| loopback_anchor_fraction = max(0.0, min(float(_setting(self.cfg, "map_loopback_anchor_fraction", 0.0)), 1.0)) |
| loopback_anchor_period = max(0, int(_setting(self.cfg, "map_loopback_anchor_period", 0))) |
| loopback_turn_gate = max(0.0, float(_setting(self.cfg, "map_loopback_turn_gate", 0.42))) |
| mid_anchor = (anchors > 1) & (anchors < route_len - 2) |
| route_turn_anchor = route_turn_strength.gather(1, anchors) |
| route_choke_anchor = mid_anchor & (route_turn_anchor > loopback_turn_gate) |
| if loopback_anchor_period > 0: |
| route_choke_anchor = route_choke_anchor & (((anchors - 1) % loopback_anchor_period) == 0) |
| mid_even = mid_anchor & even_extra |
| |
| |
| |
| mid_even_count = mid_even.sum(dim=1).clamp_min(1) |
| raw_target = (mid_even_count.to(self.device).to(torch.float32) * float(loopback_anchor_fraction)).floor().to(torch.int64) |
| if route_len <= 5 or float(loopback_anchor_fraction) <= 0.0: |
| target_loopbacks = torch.zeros((n,), device=self.device, dtype=torch.int64) |
| else: |
| target_loopbacks = raw_target.clamp_min(0).clamp_max(mid_even_count) |
| k_max = int(target_loopbacks.max().item()) if int(extras) > 0 else 0 |
| if k_max > 0: |
| choke_score = route_turn_anchor + 0.05 * self.rand(n, extras) |
| fallback_score = route_turn_anchor - 1.0 - 0.05 * self.rand(n, extras) |
| candidate_score = torch.where(route_choke_anchor, choke_score, fallback_score) |
| candidate_score = torch.where(mid_even, candidate_score, torch.full_like(candidate_score, -1.0e7)) |
| top_vals, top_idx = torch.topk(candidate_score, k=k_max, dim=1) |
| keep_by_rank = torch.arange(k_max, device=self.device)[None, :] < target_loopbacks[:, None] |
| forced_loopback_entry = torch.zeros((n, extras), device=self.device, dtype=torch.bool) |
| forced_loopback_entry = forced_loopback_entry.scatter_(1, top_idx, keep_by_rank) |
| else: |
| forced_loopback_entry = torch.zeros((n, extras), device=self.device, dtype=torch.bool) |
| forced_loopback_entry = (anchors % 2 == 0) & forced_loopback_entry |
| forced_loopback_deadend = torch.roll(forced_loopback_entry, shifts=1, dims=1) & odd_extra |
| false_summit_fraction = max(0.0, float(_setting(self.cfg, "map_false_summit_fraction", 0.0))) |
| late_anchor = anchors >= max(2, route_len - 4) |
| goal_fan_slots = extra_i >= max(0, extras - max(2, extras // 3)) |
| false_summit = (late_anchor | goal_fan_slots) & (self.rand(n, extras) < false_summit_fraction) |
| fan_width = max(1, min(4, route_len - 1)) |
| fan_target_i = (route_len - 1 - ((extra_i + (self.env_i[:, None] % fan_width)) % fan_width)).clamp_min(1) |
| fan_target_xy = gather_xy(fan_target_i) |
| false_branch_fraction = max(0.0, float(_setting(self.cfg, "map_false_branch_fraction", 0.0))) |
| false_branch = mid_anchor & (self.rand(n, extras) < false_branch_fraction) |
| false_finish_fraction = max(0.0, float(_setting(self.cfg, "map_false_finish_fraction", 0.0))) |
| late_setup_anchor = (anchors >= max(2, route_len - 7)) & (anchors < route_len - 2) |
| false_finish_entry = late_setup_anchor & even_extra & (self.rand(n, extras) < false_finish_fraction) |
| false_finish_goal = torch.roll(false_finish_entry, shifts=1, dims=1) & odd_extra |
| extra_reach_exclude = false_finish_goal |
| braid_bridge_fraction = max(0.0, float(_setting(self.cfg, "map_braid_bridge_fraction", 0.0))) |
| braid_source_ok = (anchors > 0) & (anchors < route_len - max(2, braid_span)) |
| braid_start = braid_source_ok & even_extra & (self.rand(n, extras) < braid_bridge_fraction) |
| braid_end = torch.roll(braid_start, shifts=1, dims=1) & odd_extra |
| greedy_trap_fraction = max(0.0, float(_setting(self.cfg, "map_greedy_trap_fraction", 0.0))) |
| trap_entry = mid_anchor & even_extra & ~braid_start & (self.rand(n, extras) < greedy_trap_fraction) |
| trap_deadend = torch.roll(trap_entry, shifts=1, dims=1) & odd_extra |
| false_branch = false_branch | trap_entry |
| false_branch = false_branch | forced_loopback_entry |
| trap_jump_fraction = max(0.0, float(_setting(self.cfg, "map_trap_jump_fraction", 0.0))) |
| trap_jump = mid_anchor & ~braid_start & ~braid_end & (self.rand(n, extras) < trap_jump_fraction) |
| trap_jump = trap_jump | false_finish_entry |
| trap_high_fraction = max(0.0, min(float(_setting(self.cfg, "map_trap_jump_high_fraction", 0.0)), 1.0)) |
| trap_high = trap_jump & (self.rand(n, extras) < trap_high_fraction) |
| trap_low = trap_jump & ~trap_high |
| extra_reach_exclude = extra_reach_exclude | trap_high | forced_loopback_entry | forced_loopback_deadend |
| if unlock_lever_mask is not None: |
| direct_decoy = direct_decoy & ~unlock_lever_mask |
| cluster_extra = cluster_extra & ~unlock_lever_mask |
| decoy_corridor = decoy_corridor & ~unlock_lever_mask |
| false_summit = false_summit & ~unlock_lever_mask |
| false_branch = false_branch & ~unlock_lever_mask |
| false_finish_entry = false_finish_entry & ~unlock_lever_mask |
| false_finish_goal = false_finish_goal & ~unlock_lever_mask |
| extra_reach_exclude = false_finish_goal |
| braid_start = braid_start & ~unlock_lever_mask |
| trap_entry = trap_entry & ~unlock_lever_mask |
| braid_end = torch.roll(braid_start, shifts=1, dims=1) & odd_extra & ~unlock_lever_mask |
| trap_deadend = torch.roll(trap_entry, shifts=1, dims=1) & odd_extra & ~unlock_lever_mask |
| trap_jump = trap_jump & ~unlock_lever_mask |
| trap_high = trap_high & ~unlock_lever_mask |
| trap_low = trap_jump & ~trap_high |
| extra_reach_exclude = false_finish_goal | trap_high |
| forced_loopback_entry = forced_loopback_entry & ~unlock_lever_mask |
| forced_loopback_deadend = forced_loopback_deadend & ~unlock_lever_mask |
| goal_xy = route_xy[:, -1:, :] |
| goal_dir = goal_xy - anchor_xy |
| goal_dir = goal_dir / torch.linalg.norm(goal_dir, dim=2, keepdim=True).clamp_min(1e-4) |
| direct_xy = anchor_xy + goal_dir * (step_len_base * (1.20 + 0.75 * self.rand(n, extras)))[:, :, None] |
| corridor_lane = ((extra_i + offset_seed) % 3).float() |
| corridor_xy = ( |
| anchor_xy |
| + goal_dir * (step_len_base * (0.72 + 0.42 * corridor_lane + 0.18 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.18 + 0.28 * self.rand(n, extras))[:, :, None] |
| ) |
| cluster_lane = ((extra_i + self.env_i[:, None]) % 5).float() |
| cluster_xy = ( |
| anchor_xy |
| + perp * side[:, :, None] * (island_radius * (0.78 + 0.18 * cluster_lane + 0.28 * self.rand(n, extras)))[:, :, None] |
| + tangent * ((cluster_lane - 2.0) * 0.28 * step_len_base)[:, :, None] |
| ) |
| extra_xy = torch.where(bridge_mask[:, :, None], bridge_xy, island_xy) |
| extra_xy = torch.where(terrace_mask[:, :, None], terrace_xy, extra_xy) |
| extra_xy = torch.where(cluster_extra[:, :, None], cluster_xy, extra_xy) |
| extra_xy = torch.where(decoy_corridor[:, :, None], corridor_xy, extra_xy) |
| extra_xy = torch.where(direct_decoy[:, :, None], direct_xy, extra_xy) |
| false_summit_xy = ( |
| 0.62 * fan_target_xy |
| + 0.38 * anchor_xy |
| + perp * side[:, :, None] * (0.48 + 0.24 * self.rand(n, extras))[:, :, None] |
| ) |
| false_branch_xy = ( |
| anchor_xy |
| + goal_dir * (step_len_base * (0.82 + 0.24 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.36 + 0.22 * self.rand(n, extras))[:, :, None] |
| ) |
| trap_entry_xy = ( |
| anchor_xy |
| + goal_dir * (step_len_base * (0.92 + 0.18 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.16 + 0.16 * self.rand(n, extras))[:, :, None] |
| ) |
| forced_loopback_entry_xy = ( |
| anchor_xy |
| + goal_dir * (step_len_base * (0.82 + 0.12 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.20 + 0.24 * self.rand(n, extras))[:, :, None] |
| ) |
| trap_deadend_xy = ( |
| torch.roll(trap_entry_xy, shifts=1, dims=1) |
| + torch.roll(goal_dir, shifts=1, dims=1) |
| * (step_len_base * (0.82 + 0.16 * self.rand(n, extras)))[:, :, None] |
| + torch.roll(perp * side[:, :, None], shifts=1, dims=1) |
| * (0.18 + 0.18 * self.rand(n, extras))[:, :, None] |
| ) |
| forced_loopback_deadend_xy = torch.roll(forced_loopback_entry_xy, shifts=1, dims=1) |
| trap_disguise = min(max(float(_setting(self.cfg, "map_trap_jump_disguise", 0.0)), 0.0), 1.5) |
| lure_forward = step_len_base * (0.82 + 0.34 * self.rand(n, extras)) |
| lure_side = (0.05 + (0.16 + 0.26 * trap_disguise) * self.rand(n, extras)) |
| high_forward = step_len_base * (0.58 + 0.24 * self.rand(n, extras)) |
| high_side = (0.18 + (0.30 + 0.24 * trap_disguise) * self.rand(n, extras)) |
| lateral_lure_xy = ( |
| anchor_xy |
| + tangent * (step_len_base * (0.36 + 0.26 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.74 + 0.34 * self.rand(n, extras) + 0.22 * trap_disguise)[:, :, None] |
| ) |
| direct_lure_xy = ( |
| anchor_xy |
| + goal_dir * lure_forward[:, :, None] |
| + perp * side[:, :, None] * lure_side[:, :, None] |
| ) |
| high_lure_xy = ( |
| anchor_xy |
| + goal_dir * high_forward[:, :, None] |
| + perp * side[:, :, None] * high_side[:, :, None] |
| ) |
| disguised_trap = self.rand(n, extras) < min(0.70, 0.34 * trap_disguise) |
| trap_jump_xy = torch.where(disguised_trap[:, :, None], lateral_lure_xy, direct_lure_xy) |
| trap_jump_xy = torch.where(trap_high[:, :, None], high_lure_xy, trap_jump_xy) |
| false_finish_entry_xy = ( |
| anchor_xy |
| + goal_dir * (step_len_base * (0.86 + 0.18 * self.rand(n, extras)))[:, :, None] |
| + perp * side[:, :, None] * (0.12 + 0.16 * self.rand(n, extras))[:, :, None] |
| ) |
| finish_dir = torch.roll(goal_dir, shifts=1, dims=1) |
| finish_perp = torch.roll(perp, shifts=1, dims=1) |
| finish_side = torch.roll(side, shifts=1, dims=1) |
| false_finish_goal_xy = ( |
| goal_xy |
| - finish_dir * (step_len_base * (0.84 + 0.18 * self.rand(n, extras)))[:, :, None] |
| + finish_perp * finish_side[:, :, None] * (0.30 + 0.30 * self.rand(n, extras))[:, :, None] |
| ) |
| bridge_vec = braid_target_xy - anchor_xy |
| bridge_dir = bridge_vec / torch.linalg.norm(bridge_vec, dim=2, keepdim=True).clamp_min(1e-4) |
| bridge_perp = torch.stack((-bridge_dir[..., 1], bridge_dir[..., 0]), dim=2) |
| bridge_side = torch.where(((self.env_i[:, None] + anchors) % 2) == 0, 1.0, -1.0) |
| bridge_bulge = island_radius * (0.54 + 0.18 * self.rand(n, extras)) |
| braid_start_xy = ( |
| anchor_xy * (1.0 - 0.34) |
| + braid_target_xy * 0.34 |
| + bridge_perp * bridge_side[:, :, None] * bridge_bulge[:, :, None] |
| ) |
| braid_end_xy = ( |
| torch.roll(anchor_xy, shifts=1, dims=1) * (1.0 - 0.68) |
| + torch.roll(braid_target_xy, shifts=1, dims=1) * 0.68 |
| + torch.roll(bridge_perp * bridge_side[:, :, None] * bridge_bulge[:, :, None] * 0.78, shifts=1, dims=1) |
| ) |
| extra_xy = torch.where(false_branch[:, :, None], false_branch_xy, extra_xy) |
| extra_xy = torch.where(trap_entry[:, :, None], trap_entry_xy, extra_xy) |
| extra_xy = torch.where(trap_deadend[:, :, None], trap_deadend_xy, extra_xy) |
| extra_xy = torch.where(forced_loopback_entry[:, :, None], forced_loopback_entry_xy, extra_xy) |
| extra_xy = torch.where(forced_loopback_deadend[:, :, None], forced_loopback_deadend_xy, extra_xy) |
| extra_xy = torch.where(braid_start[:, :, None], braid_start_xy, extra_xy) |
| extra_xy = torch.where(braid_end[:, :, None], braid_end_xy, extra_xy) |
| extra_xy = torch.where(false_summit[:, :, None], false_summit_xy, extra_xy) |
| extra_xy = torch.where(trap_jump[:, :, None], trap_jump_xy, extra_xy) |
| extra_xy = torch.where(false_finish_entry[:, :, None], false_finish_entry_xy, extra_xy) |
| extra_xy = torch.where(false_finish_goal[:, :, None], false_finish_goal_xy, extra_xy) |
| extra_xy = clamp_xy_to_route_span( |
| route_xy, |
| extra_xy, |
| min_frac=float(_setting(self.cfg, "map_extra_span_min_frac", 0.05)), |
| max_frac=float(_setting(self.cfg, "map_extra_span_max_frac", 0.90)), |
| end_margin=float(_setting(self.cfg, "map_extra_span_end_margin", 1.10)), |
| ) |
| if unlock_lever_mask is not None: |
| unlock_dir = next_xy - anchor_xy |
| unlock_dir = unlock_dir / torch.linalg.norm(unlock_dir, dim=2, keepdim=True).clamp_min(1e-4) |
| unlock_perp = torch.stack((-unlock_dir[..., 1], unlock_dir[..., 0]), dim=2) |
| unlock_side = torch.where(((self.env_i[:, None] + anchors) % 2) == 0, 1.0, -1.0) |
| unlock_xy = ( |
| anchor_xy |
| + unlock_dir * (step_len_base * (0.40 + 0.14 * self.rand(n, extras)))[:, :, None] |
| + unlock_perp * unlock_side[:, :, None] * (island_radius * (0.92 + 0.18 * self.rand(n, extras)))[:, :, None] |
| ) |
| extra_xy = torch.where(unlock_lever_mask[:, :, None], unlock_xy, extra_xy) |
| self.platforms[:, route_len:, :2] = extra_xy |
|
|
| anchor_z = route_z.gather(1, anchors) |
| braid_target_z = route_z.gather(1, braid_target_i) |
| fan_target_z = route_z.gather(1, fan_target_i) |
| legal_z = anchor_z + (self.rand(n, extras) - 0.5) * vertical_step * 0.90 |
| legal_z = torch.where(terrace_mask, anchor_z + vertical_step * (0.38 + 0.20 * self.rand(n, extras)), legal_z) |
| cluster_z = anchor_z + (self.rand(n, extras) - 0.5) * vertical_step * (1.15 + 0.35 * hairpin_depth) |
| legal_z = torch.where(cluster_extra, cluster_z, legal_z) |
| branch_z = anchor_z + (self.rand(n, extras) - 0.35) * vertical_step * 0.34 |
| legal_z = torch.where(false_branch, branch_z, legal_z) |
| summit_legal_z = anchor_z * 0.38 + fan_target_z * 0.62 + (self.rand(n, extras) - 0.45) * vertical_step * 0.20 |
| legal_z = torch.where(false_summit, summit_legal_z, legal_z) |
| trap_deadend_z = torch.roll(branch_z, shifts=1, dims=1) + (self.rand(n, extras) - 0.45) * vertical_step * 0.28 |
| legal_z = torch.where(trap_deadend, trap_deadend_z, legal_z) |
| braid_start_z = anchor_z * (1.0 - 0.34) + braid_target_z * 0.34 - vertical_step * (0.10 + 0.10 * self.rand(n, extras)) |
| braid_end_z = ( |
| torch.roll(anchor_z, shifts=1, dims=1) * (1.0 - 0.68) |
| + torch.roll(braid_target_z, shifts=1, dims=1) * 0.68 |
| + vertical_step * (0.04 + 0.10 * self.rand(n, extras)) |
| ) |
| legal_z = torch.where(braid_start, braid_start_z, legal_z) |
| legal_z = torch.where(braid_end, braid_end_z, legal_z) |
| goal_z = route_z[:, -1:] |
| false_finish_goal_z = goal_z - vertical_step * (0.08 + 0.10 * self.rand(n, extras)) |
| legal_z = torch.where(false_finish_goal, false_finish_goal_z, legal_z) |
| trap_drop_max = max(0.55, float(_setting(self.cfg, "map_trap_jump_drop", 2.2))) |
| trap_drop_min = min(trap_drop_max - 0.04, max(0.12, float(_setting(self.cfg, "map_trap_jump_min_drop", 0.42)))) |
| trap_edge_margin = max(0.0, float(_setting(self.cfg, "map_trap_jump_edge_margin", 0.10))) |
| trap_deep_z = anchor_z - trap_drop_max |
| trap_shallow_z = anchor_z - trap_drop_min |
| source_size_for_trap = self.sizes[:, :route_len].gather(1, anchors[:, :, None].expand(-1, -1, 2)) |
| trap_size = self.sizes[:, route_len:] |
| anchor_pos_for_trap = torch.cat((anchor_xy, anchor_z[:, :, None]), dim=2) |
| lo = trap_deep_z |
| hi = trap_shallow_z |
| for _ in range(7): |
| mid = (lo + hi) * 0.5 |
| trap_pos = torch.cat((trap_jump_xy, mid[:, :, None]), dim=2) |
| return_margin = pair_reach_margin(self, trap_pos, anchor_pos_for_trap, trap_size, source_size_for_trap) |
| too_recoverable = return_margin > -trap_edge_margin |
| hi = torch.where(too_recoverable, mid, hi) |
| lo = torch.where(too_recoverable, lo, mid) |
| trap_jump_z = lo + (self.rand(n, extras) - 0.5) * 0.04 |
| high_min_up = max(0.08, float(_setting(self.cfg, "map_trap_jump_high_min_up", 0.34))) |
| high_max_up = max(high_min_up + 0.08, float(_setting(self.cfg, "map_trap_jump_high_max_up", 1.35))) |
| high_lo = anchor_z + high_min_up |
| high_hi = torch.minimum(anchor_z + high_max_up, torch.full_like(anchor_z, max_h)) |
| for _ in range(8): |
| mid = (high_lo + high_hi) * 0.5 |
| trap_pos = torch.cat((trap_jump_xy, mid[:, :, None]), dim=2) |
| approach_margin = pair_reach_margin(self, anchor_pos_for_trap, trap_pos, source_size_for_trap, trap_size) |
| too_reachable = approach_margin > -trap_edge_margin |
| high_lo = torch.where(too_reachable, mid, high_lo) |
| high_hi = torch.where(too_reachable, high_hi, mid) |
| trap_high_z = high_hi + (0.03 + 0.05 * self.rand(n, extras)) |
| trap_high_z = torch.minimum(trap_high_z, torch.full_like(trap_high_z, max_h)) |
| trap_jump_z = torch.where(trap_high, trap_high_z, trap_jump_z) |
| legal_z = torch.where(trap_jump, trap_jump_z, legal_z) |
| if unlock_lever_mask is not None: |
| unlock_next_z = route_z.gather(1, (anchors + 1).clamp_max(route_len - 1)) |
| unlock_z = anchor_z * 0.56 + unlock_next_z * 0.44 + (self.rand(n, extras) - 0.5) * vertical_step * 0.14 |
| legal_z = torch.where(unlock_lever_mask, unlock_z, legal_z) |
| self.sizes[:, route_len : route_len + unlock_gate_count] = torch.maximum( |
| self.sizes[:, route_len : route_len + unlock_gate_count], |
| base_size * 1.02, |
| ) |
| decoy = self.rand(n, extras) < float(getattr(self.cfg, "map_illegal_decoy_fraction", 0.0)) |
| if unlock_lever_mask is not None: |
| decoy = decoy & ~unlock_lever_mask |
| |
| |
| illegal_z = anchor_z + float(getattr(self.cfg, "map_illegal_decoy_min_up", 0.9)) + self.rand(n, extras) * ( |
| float(getattr(self.cfg, "map_illegal_decoy_max_up", 1.45)) - float(getattr(self.cfg, "map_illegal_decoy_min_up", 0.9)) |
| ) |
| extra_z = torch.where(decoy, illegal_z, legal_z) |
| normal_floor = torch.full_like(extra_z, height_offset + 0.18) |
| trap_floor = torch.full_like(extra_z, float(self.lava_z) + 0.65) |
| z_floor = torch.where(trap_low, trap_floor, normal_floor) |
| self.platforms[:, route_len:, 2] = torch.minimum(torch.maximum(extra_z, z_floor), torch.full_like(extra_z, max_h)) |
| source_pos = torch.cat( |
| ( |
| anchor_xy, |
| anchor_z[:, :, None], |
| ), |
| dim=2, |
| ) |
| source_size = self.sizes[:, :route_len].gather(1, anchors[:, :, None].expand(-1, -1, 2)) |
| rolled_extra_pos = torch.roll(self.platforms[:, route_len:], shifts=1, dims=1) |
| rolled_extra_size = torch.roll(self.sizes[:, route_len:], shifts=1, dims=1) |
| source_pos = torch.where(trap_deadend[:, :, None], rolled_extra_pos, source_pos) |
| source_size = torch.where(trap_deadend[:, :, None], rolled_extra_size, source_size) |
| source_pos = torch.where(forced_loopback_deadend[:, :, None], rolled_extra_pos, source_pos) |
| source_size = torch.where(forced_loopback_deadend[:, :, None], rolled_extra_size, source_size) |
| source_pos = torch.where(braid_end[:, :, None], rolled_extra_pos, source_pos) |
| source_size = torch.where(braid_end[:, :, None], rolled_extra_size, source_size) |
| self_source_pos = self.platforms[:, route_len:] |
| self_source_size = self.sizes[:, route_len:] |
| source_pos = torch.where(false_finish_goal[:, :, None], self_source_pos, source_pos) |
| source_size = torch.where(false_finish_goal[:, :, None], self_source_size, source_size) |
| enforce_extra_source_reach(self, source_pos, source_size, safety=0.10, passes=1, exclude_mask=trap_high) |
| if unlock_lever_mask is not None and unlock_target_i is not None: |
| enforce_masked_extra_to_route_reach(self, unlock_lever_mask, unlock_target_i, safety=0.08, passes=2) |
| if extra_reach_exclude is not None: |
| goal_index = torch.full_like(anchors, route_len - 1) |
| enforce_masked_extra_to_route_reach(self, extra_reach_exclude, goal_index, safety=0.08, passes=2) |
| enforce_masked_extra_to_route_reach(self, braid_end, torch.roll(braid_target_i, shifts=1, dims=1), safety=0.04, passes=1) |
| enforce_masked_extra_to_route_reach(self, false_summit, fan_target_i, safety=0.04, passes=1) |
| enforce_extra_source_reach(self, source_pos, source_size, safety=0.04, passes=1, exclude_mask=trap_high) |
|
|
| self.separate_xy_overlaps(iterations=cleanup_iters) |
| separate_route_overlaps(self, iterations=max(12, cleanup_iters // 8), clearance=0.10) |
| separate_nonroute_overlaps(self, iterations=max(24, cleanup_iters // 2)) |
| enforce_route_edge_gap(self, floor=max(0.0, float(_setting(self.cfg, "map_route_edge_gap_floor", 0.0))), passes=2) |
| enforce_route_reach(self, safety=0.12, passes=2) |
| relocate_route_colliding_extras(self, iterations=5, clearance=0.26) |
| separate_nonroute_overlaps(self, iterations=max(32, cleanup_iters)) |
| enforce_route_edge_gap(self, floor=max(0.0, float(_setting(self.cfg, "map_route_edge_gap_floor", 0.0))), passes=2) |
| enforce_route_reach(self, safety=0.12, passes=2) |
| relocate_route_colliding_extras(self, iterations=6, clearance=0.30) |
| resolve_route_extra_grazes(self, iterations=8, clearance=0.06) |
| scatter_colliding_extras(self, iterations=8, clearance=0.34) |
| relocate_route_colliding_extras(self, iterations=4, clearance=0.34) |
| resolve_route_extra_grazes(self, iterations=4, clearance=0.08) |
| separate_route_overlaps(self, iterations=max(8, cleanup_iters // 16), clearance=0.08) |
| enforce_route_edge_gap(self, floor=max(0.0, float(_setting(self.cfg, "map_route_edge_gap_floor", 0.0))), passes=2) |
| enforce_route_reach(self, safety=0.12, passes=2) |
| separate_route_overlaps(self, iterations=max(16, cleanup_iters // 8), clearance=0.08) |
| enforce_route_edge_gap(self, floor=max(0.0, float(_setting(self.cfg, "map_route_edge_gap_floor", 0.0))), passes=1) |
| enforce_route_reach(self, safety=0.08, passes=1) |
| separate_route_overlaps(self, iterations=max(8, cleanup_iters // 12), clearance=0.04) |
| enforce_route_reach(self, safety=0.12, passes=2) |
| enforce_extra_any_route_reach(self, safety=0.02, passes=2, exclude_mask=extra_reach_exclude) |
| if int(self.p) > int(self.route_len): |
| self.platforms[:, self.route_len :, :2] = clamp_xy_to_route_span( |
| self.platforms[:, : self.route_len, :2], |
| self.platforms[:, self.route_len :, :2], |
| min_frac=float(_setting(self.cfg, "map_extra_span_min_frac", 0.05)), |
| max_frac=float(_setting(self.cfg, "map_extra_span_max_frac", 0.90)), |
| end_margin=float(_setting(self.cfg, "map_extra_span_end_margin", 1.10)), |
| ) |
| resolve_route_extra_grazes(self, iterations=4, clearance=0.05) |
| enforce_extra_any_route_reach(self, safety=0.00, passes=1, exclude_mask=extra_reach_exclude) |
| if int(self.p) > int(self.route_len): |
| self.platforms[:, self.route_len :, :2] = clamp_xy_to_route_span( |
| self.platforms[:, : self.route_len, :2], |
| self.platforms[:, self.route_len :, :2], |
| min_frac=float(_setting(self.cfg, "map_extra_span_min_frac", 0.05)), |
| max_frac=float(_setting(self.cfg, "map_extra_span_max_frac", 0.90)), |
| end_margin=float(_setting(self.cfg, "map_extra_span_end_margin", 1.10)), |
| ) |
| center = self.platforms[:, :, :2].mean(dim=1, keepdim=True) |
| self.platforms[:, :, :2] -= center |
|
|
|
|
| def clamp_xy_to_route_span( |
| route_xy: torch.Tensor, |
| xy: torch.Tensor, |
| *, |
| min_frac: float = 0.05, |
| max_frac: float = 0.90, |
| end_margin: float = 1.10, |
| ) -> torch.Tensor: |
| if xy.numel() == 0 or route_xy.shape[1] < 2: |
| return xy |
| start = route_xy[:, 0:1] |
| goal = route_xy[:, -1:] |
| axis = goal - start |
| span = torch.linalg.norm(axis, dim=2, keepdim=True).clamp_min(1e-4) |
| axis = axis / span |
| rel = xy - start |
| proj = (rel * axis).sum(dim=2, keepdim=True) |
| lower = span * float(min_frac) |
| upper = torch.maximum(lower + 0.25, span * float(max_frac) - float(end_margin)) |
| clamped = proj.clamp_min(lower).clamp_max(upper) |
| return xy + axis * (clamped - proj) |
|
|
|
|
| def overlap_count(platforms: torch.Tensor, sizes: torch.Tensor, *, margin: float = 0.0) -> int: |
| p = int(platforms.shape[1]) |
| xy = platforms[:, :, :2] |
| delta = (xy[:, :, None, :] - xy[:, None, :, :]).abs() |
| gap = (sizes[:, :, None, :] + sizes[:, None, :, :]) * 0.5 + float(margin) |
| colliding = (delta[..., 0] < gap[..., 0]) & (delta[..., 1] < gap[..., 1]) |
| eye = torch.eye(p, dtype=torch.bool, device=platforms.device)[None, :, :] |
| return int(((colliding & ~eye).sum() // 2).detach().cpu()) |
|
|
|
|
| def route_overlap_count(env: Any, *, margin: float = 0.0) -> int: |
| return overlap_count(env.platforms[:, : env.route_len], env.sizes[:, : env.route_len], margin=margin) |
|
|
|
|
| def extra_overlap_count(env: Any, *, margin: float = 0.0) -> int: |
| route_len = int(env.route_len) |
| if int(env.p) <= route_len: |
| return 0 |
| return overlap_count(env.platforms[:, route_len:], env.sizes[:, route_len:], margin=margin) |
|
|
|
|
| def route_extra_overlap_count(env: Any, *, margin: float = 0.0) -> int: |
| route_len, p = int(env.route_len), int(env.p) |
| if p <= route_len: |
| return 0 |
| route_xy = env.platforms[:, :route_len, :2] |
| extra_xy = env.platforms[:, route_len:, :2] |
| route_sizes = env.sizes[:, :route_len] |
| extra_sizes = env.sizes[:, route_len:] |
| delta = (extra_xy[:, :, None, :] - route_xy[:, None, :, :]).abs() |
| gap = (extra_sizes[:, :, None, :] + route_sizes[:, None, :, :]) * 0.5 + float(margin) |
| colliding = (delta[..., 0] < gap[..., 0]) & (delta[..., 1] < gap[..., 1]) |
| return int(colliding.sum().detach().cpu()) |
|
|
|
|
| def pair_reach_margin(env: Any, src: torch.Tensor, dst: torch.Tensor, src_size: torch.Tensor, dst_size: torch.Tensor) -> torch.Tensor: |
| delta = dst - src |
| xy = torch.linalg.norm(delta[..., :2], dim=-1) |
| dz = delta[..., 2] |
| jump_v = float(env.jump_velocity) |
| gravity = float(env.gravity) |
| disc = jump_v * jump_v - 2.0 * gravity * dz |
| feasible = disc >= 0.0 |
| flight = (jump_v + torch.sqrt(disc.clamp_min(0.0))) / gravity |
| reach = float(env.max_speed) * float(env.sprint_speed_mult) * flight * float(getattr(env.cfg, "map_reach_scale", 0.82)) |
| landing_pad = dst_size.amin(dim=-1) * 0.35 + src_size.amin(dim=-1) * 0.18 + float(env.radius) * 0.8 |
| allowed = torch.where(feasible, reach + landing_pad - 0.10, torch.full_like(reach, -1.0e6)) |
| return allowed - xy |
|
|
|
|
| def enforce_extra_source_reach( |
| env: Any, |
| source_pos: torch.Tensor, |
| source_size: torch.Tensor, |
| *, |
| safety: float = 0.10, |
| passes: int = 2, |
| exclude_mask: torch.Tensor | None = None, |
| ) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| for _ in range(max(1, int(passes))): |
| dst = env.platforms[:, route_len:] |
| dst_size = env.sizes[:, route_len:] |
| margin = pair_reach_margin(env, source_pos, dst, source_size, dst_size) |
| too_far = margin < float(safety) |
| if exclude_mask is not None: |
| too_far = too_far & ~exclude_mask |
| if not bool(too_far.any().detach().cpu()): |
| return |
| delta_xy = dst[:, :, :2] - source_pos[:, :, :2] |
| dist = torch.linalg.norm(delta_xy, dim=2).clamp_min(1e-4) |
| direction = delta_xy / dist[:, :, None] |
| pull = (float(safety) - margin).clamp_min(0.0) |
| env.platforms[:, route_len:, :2] = dst[:, :, :2] - direction * pull[:, :, None] * too_far[:, :, None].float() |
|
|
|
|
| def enforce_extra_any_route_reach( |
| env: Any, |
| *, |
| safety: float = 0.0, |
| passes: int = 2, |
| exclude_mask: torch.Tensor | None = None, |
| ) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| for _ in range(max(1, int(passes))): |
| route_pos = env.platforms[:, :route_len] |
| route_size = env.sizes[:, :route_len] |
| extra_pos = env.platforms[:, route_len:] |
| extra_size = env.sizes[:, route_len:] |
| margins = pair_reach_margin( |
| env, |
| route_pos[:, :, None, :], |
| extra_pos[:, None, :, :], |
| route_size[:, :, None, :], |
| extra_size[:, None, :, :], |
| ) |
| best_margin, best_route = margins.max(dim=1) |
| too_far = best_margin < float(safety) |
| if exclude_mask is not None: |
| too_far = too_far & ~exclude_mask |
| if not bool(too_far.any().detach().cpu()): |
| return |
| src = route_pos.gather(1, best_route[:, :, None].expand(-1, -1, 3)) |
| delta_xy = extra_pos[:, :, :2] - src[:, :, :2] |
| dist = torch.linalg.norm(delta_xy, dim=2).clamp_min(1e-4) |
| direction = delta_xy / dist[:, :, None] |
| pull = (float(safety) - best_margin).clamp_min(0.0) |
| env.platforms[:, route_len:, :2] = extra_pos[:, :, :2] - direction * pull[:, :, None] * too_far[:, :, None].float() |
|
|
|
|
| def enforce_masked_extra_to_route_reach( |
| env: Any, |
| mask: torch.Tensor, |
| route_index: torch.Tensor, |
| *, |
| safety: float = 0.0, |
| passes: int = 1, |
| ) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| route_index = route_index.clamp(0, route_len - 1) |
| for _ in range(max(1, int(passes))): |
| extra_pos = env.platforms[:, route_len:] |
| extra_size = env.sizes[:, route_len:] |
| dst = env.platforms[:, :route_len].gather(1, route_index[:, :, None].expand(-1, -1, 3)) |
| dst_size = env.sizes[:, :route_len].gather(1, route_index[:, :, None].expand(-1, -1, 2)) |
| margin = pair_reach_margin(env, extra_pos, dst, extra_size, dst_size) |
| too_far = (margin < float(safety)) & mask |
| if not bool(too_far.any().detach().cpu()): |
| return |
| delta_xy = extra_pos[:, :, :2] - dst[:, :, :2] |
| dist = torch.linalg.norm(delta_xy, dim=2).clamp_min(1e-4) |
| direction = delta_xy / dist[:, :, None] |
| pull = (float(safety) - margin).clamp_min(0.0) |
| env.platforms[:, route_len:, :2] = extra_pos[:, :, :2] - direction * pull[:, :, None] * too_far[:, :, None].float() |
|
|
|
|
| def route_reach_margin(env: Any) -> torch.Tensor: |
| route = env.platforms[:, : env.route_len] |
| sizes = env.sizes[:, : env.route_len] |
| delta = route[:, 1:] - route[:, :-1] |
| xy = torch.linalg.norm(delta[:, :, :2], dim=2) |
| dz = delta[:, :, 2] |
| jump_v = float(env.jump_velocity) |
| gravity = float(env.gravity) |
| disc = jump_v * jump_v - 2.0 * gravity * dz |
| feasible = disc >= 0.0 |
| flight = (jump_v + torch.sqrt(disc.clamp_min(0.0))) / gravity |
| reach = float(env.max_speed) * float(env.sprint_speed_mult) * flight * float(getattr(env.cfg, "map_reach_scale", 0.82)) |
| landing_pad = sizes[:, 1:].amin(dim=2) * 0.35 + sizes[:, :-1].amin(dim=2) * 0.18 + float(env.radius) * 0.8 |
| allowed = torch.where(feasible, reach + landing_pad - 0.10, torch.full_like(reach, -1.0e6)) |
| return allowed - xy |
|
|
|
|
| def all_platform_path_diagnostics(env: Any, *, reach_margin: float = 0.0) -> dict[str, float]: |
| """No-route-label geometry audit over every platform. |
| |
| This is intentionally separate from the construction route. It treats every |
| physically reachable platform transition as a graph edge, then measures the |
| shortest start-to-goal path a perfect full-map planner could take. If this is |
| much shorter than route_len - 1, the generated map is easier than its route |
| count claims. |
| """ |
| n, p = int(env.n), int(env.p) |
| route_jumps = max(1, int(env.route_len) - 1) |
| if p <= 1: |
| return { |
| "all_path_goal_reachable_fraction": 0.0, |
| "all_path_shortest_hops_mean": float(route_jumps + 1), |
| "all_path_shortest_hops_min": float(route_jumps + 1), |
| "all_path_shortest_hops_p50": float(route_jumps + 1), |
| "all_path_shortcut_ratio_mean": float((route_jumps + 1) / route_jumps), |
| "all_path_shortcut_le_half_fraction": 0.0, |
| "all_path_shortcut_le_five_fraction": 0.0, |
| "all_path_platform_reachable_fraction": 0.0, |
| "all_platform_nearest_edge_gap_mean": 0.0, |
| "all_platform_nearest_edge_gap_min": 0.0, |
| } |
|
|
| pos = env.platforms |
| sizes = env.sizes |
| margin = pair_reach_margin( |
| env, |
| pos[:, :, None, :], |
| pos[:, None, :, :], |
| sizes[:, :, None, :], |
| sizes[:, None, :, :], |
| ) |
| eye = torch.eye(p, dtype=torch.bool, device=env.device)[None, :, :] |
| edge = (margin > float(reach_margin)) & ~eye |
| reached = torch.zeros((n, p), dtype=torch.bool, device=env.device) |
| reached[:, 0] = True |
| distance = torch.full((n, p), float(p + 1), device=env.device) |
| distance[:, 0] = 0.0 |
| for hop in range(1, p + 1): |
| next_reached = (reached[:, :, None] & edge).any(dim=1) |
| new = next_reached & ~reached |
| distance = torch.where(new, torch.full_like(distance, float(hop)), distance) |
| reached = reached | next_reached |
| if bool(reached[:, int(env.route_len) - 1].all().detach().cpu()): |
| break |
|
|
| goal_hops = distance[:, int(env.route_len) - 1] |
| goal_reachable = goal_hops <= float(p) |
| ratio = goal_hops / float(route_jumps) |
| half_threshold = max(1.0, float(route_jumps) * 0.50) |
|
|
| delta_xy = pos[:, :, None, :2] - pos[:, None, :, :2] |
| dist_xy = torch.linalg.norm(delta_xy, dim=3).clamp_min(1e-4) |
| direction = delta_xy / dist_xy[:, :, :, None] |
| support = ( |
| (direction.abs() * sizes[:, :, None, :]).sum(dim=3) * 0.5 |
| + (direction.abs() * sizes[:, None, :, :]).sum(dim=3) * 0.5 |
| ) |
| edge_gap = dist_xy - support |
| edge_gap = torch.where(eye, torch.full_like(edge_gap, 1.0e6), edge_gap) |
| nearest_gap = edge_gap.amin(dim=2) |
|
|
| return { |
| "all_path_goal_reachable_fraction": float(goal_reachable.float().mean().detach().cpu()), |
| "all_path_shortest_hops_mean": float(goal_hops.mean().detach().cpu()), |
| "all_path_shortest_hops_min": float(goal_hops.min().detach().cpu()), |
| "all_path_shortest_hops_p50": float(torch.quantile(goal_hops, 0.50).detach().cpu()), |
| "all_path_shortcut_ratio_mean": float(ratio.mean().detach().cpu()), |
| "all_path_shortcut_le_half_fraction": float(((goal_hops <= half_threshold) & goal_reachable).float().mean().detach().cpu()), |
| "all_path_shortcut_le_five_fraction": float(((goal_hops <= 5.0) & goal_reachable).float().mean().detach().cpu()), |
| "all_path_platform_reachable_fraction": float(reached.float().mean().detach().cpu()), |
| "all_platform_nearest_edge_gap_mean": float(nearest_gap.mean().detach().cpu()), |
| "all_platform_nearest_edge_gap_min": float(nearest_gap.min().detach().cpu()), |
| } |
|
|
|
|
| def enforce_route_reach(env: Any, *, safety: float = 0.10, passes: int = 2) -> None: |
| route_len = int(env.route_len) |
| if route_len <= 1: |
| return |
| for _ in range(max(1, int(passes))): |
| for layer in range(1, route_len): |
| prev = env.platforms[:, layer - 1] |
| cur = env.platforms[:, layer] |
| dxy = cur[:, :2] - prev[:, :2] |
| dist = torch.linalg.norm(dxy, dim=1).clamp_min(1e-4) |
| direction = dxy / dist[:, None] |
| dz = cur[:, 2] - prev[:, 2] |
| jump_v = float(env.jump_velocity) |
| gravity = float(env.gravity) |
| disc = jump_v * jump_v - 2.0 * gravity * dz |
| flight = (jump_v + torch.sqrt(disc.clamp_min(0.0))) / gravity |
| sizes = env.sizes |
| landing_pad = sizes[:, layer].amin(dim=1) * 0.35 + sizes[:, layer - 1].amin(dim=1) * 0.18 + float(env.radius) * 0.8 |
| max_dist = ( |
| float(env.max_speed) |
| * float(env.sprint_speed_mult) |
| * flight |
| * float(getattr(env.cfg, "map_reach_scale", 0.82)) |
| + landing_pad |
| - 0.10 |
| - float(safety) |
| ) |
| too_far = dist > max_dist |
| if not bool(too_far.any().detach().cpu()): |
| continue |
| pull = direction * (dist - max_dist).clamp_min(0.0)[:, None] * too_far[:, None].float() |
| env.platforms[:, layer:route_len, :2] -= pull[:, None, :] |
|
|
|
|
| def enforce_route_edge_gap(env: Any, *, floor: float = 0.0, passes: int = 2) -> None: |
| route_len = int(env.route_len) |
| if route_len <= 1 or floor <= 0.0: |
| return |
| for _ in range(max(1, int(passes))): |
| for layer in range(1, route_len): |
| prev = env.platforms[:, layer - 1] |
| cur = env.platforms[:, layer] |
| dxy = cur[:, :2] - prev[:, :2] |
| dist = torch.linalg.norm(dxy, dim=1).clamp_min(1e-4) |
| direction = dxy / dist[:, None] |
| support = ( |
| (direction.abs() * env.sizes[:, layer - 1]).sum(dim=1) * 0.5 |
| + (direction.abs() * env.sizes[:, layer]).sum(dim=1) * 0.5 |
| ) |
| edge_gap = dist - support |
| deficit = (float(floor) - edge_gap).clamp_min(0.0) |
| env.platforms[:, layer:route_len, :2] += direction[:, None, :] * deficit[:, None, None] |
|
|
|
|
| def separate_nonroute_overlaps(env: Any, *, iterations: int = 24) -> None: |
| p = int(env.p) |
| if p <= 1: |
| return |
| margin = max(0.12, float(getattr(env.cfg, "map_braid_mesh_clearance", 0.16))) |
| eye = torch.eye(p, dtype=torch.bool, device=env.device)[None, :, :] |
| mobility = torch.ones((env.n, p), device=env.device) |
| mobility[:, : int(env.route_len)] = 0.0 |
| for _ in range(max(0, int(iterations))): |
| xy = env.platforms[:, :, :2] |
| delta = xy[:, :, None, :] - xy[:, None, :, :] |
| abs_delta = delta.abs() |
| gap_x = (env.sizes[:, :, None, 0] + env.sizes[:, None, :, 0]) * 0.5 + margin |
| gap_y = (env.sizes[:, :, None, 1] + env.sizes[:, None, :, 1]) * 0.5 + margin |
| overlap_x = gap_x - abs_delta[..., 0] |
| overlap_y = gap_y - abs_delta[..., 1] |
| colliding = (overlap_x > 0.0) & (overlap_y > 0.0) & ~eye |
| use_x = overlap_x <= overlap_y |
| sign_x = torch.where(delta[..., 0] >= 0.0, 1.0, -1.0) |
| sign_y = torch.where(delta[..., 1] >= 0.0, 1.0, -1.0) |
| pair_push = torch.zeros_like(delta) |
| pair_push[..., 0] = torch.where(colliding & use_x, sign_x * overlap_x, torch.zeros_like(overlap_x)) |
| pair_push[..., 1] = torch.where(colliding & ~use_x, sign_y * overlap_y, torch.zeros_like(overlap_y)) |
| share = mobility[:, :, None] / (mobility[:, :, None] + mobility[:, None, :] + 1e-6) |
| weighted = pair_push * share[..., None] |
| active_pairs = colliding.float().sum(dim=2, keepdim=True).clamp_min(1.0) |
| push = weighted.sum(dim=2) / active_pairs |
| env.platforms[:, :, :2] = xy + (push * mobility[:, :, None]).clamp(-0.28, 0.28) |
|
|
|
|
| def separate_route_overlaps(env: Any, *, iterations: int = 6, clearance: float = 0.06) -> None: |
| route_len = int(env.route_len) |
| if route_len <= 2: |
| return |
| ids = torch.arange(route_len, device=env.device) |
| pair_index_gap = (ids[:, None] - ids[None, :]).abs() |
| non_adjacent = pair_index_gap > 1 |
| eye = torch.eye(route_len, dtype=torch.bool, device=env.device)[None, :, :] |
| mobility = torch.linspace(0.15, 1.0, route_len, device=env.device)[None, :] |
| mobility[:, 0] = 0.02 |
| mobility[:, -1] = 0.35 |
| for _ in range(max(1, int(iterations))): |
| xy = env.platforms[:, :route_len, :2] |
| sizes = env.sizes[:, :route_len] |
| delta = xy[:, :, None, :] - xy[:, None, :, :] |
| abs_delta = delta.abs() |
| gap_x = (sizes[:, :, None, 0] + sizes[:, None, :, 0]) * 0.5 + float(clearance) |
| gap_y = (sizes[:, :, None, 1] + sizes[:, None, :, 1]) * 0.5 + float(clearance) |
| overlap_x = gap_x - abs_delta[..., 0] |
| overlap_y = gap_y - abs_delta[..., 1] |
| colliding = (overlap_x > 0.0) & (overlap_y > 0.0) & non_adjacent[None, :, :] & ~eye |
| if not bool(colliding.any().detach().cpu()): |
| return |
| use_x = overlap_x <= overlap_y |
| sign_x = torch.where(delta[..., 0] >= 0.0, 1.0, -1.0) |
| sign_y = torch.where(delta[..., 1] >= 0.0, 1.0, -1.0) |
| pair_push = torch.zeros_like(delta) |
| pair_push[..., 0] = torch.where(colliding & use_x, sign_x * overlap_x, torch.zeros_like(overlap_x)) |
| pair_push[..., 1] = torch.where(colliding & ~use_x, sign_y * overlap_y, torch.zeros_like(overlap_y)) |
| share = mobility[:, :, None] / (mobility[:, :, None] + mobility[:, None, :] + 1e-6) |
| active_pairs = colliding.float().sum(dim=2, keepdim=True).clamp_min(1.0) |
| push = (pair_push * share[..., None]).sum(dim=2) / active_pairs |
| env.platforms[:, :route_len, :2] = xy + (push * mobility[:, :, None]).clamp(-0.36, 0.36) |
|
|
|
|
| def relocate_route_colliding_extras(env: Any, *, iterations: int = 4, clearance: float = 0.22) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| route_ids = torch.arange(route_len, device=env.device) |
| extra_ids = torch.arange(extras, device=env.device) |
| for step in range(max(1, int(iterations))): |
| route_xy = env.platforms[:, :route_len, :2] |
| extra_xy = env.platforms[:, route_len:, :2] |
| route_sizes = env.sizes[:, :route_len] |
| extra_sizes = env.sizes[:, route_len:] |
| delta = extra_xy[:, :, None, :] - route_xy[:, None, :, :] |
| abs_delta = delta.abs() |
| gap = (extra_sizes[:, :, None, :] + route_sizes[:, None, :, :]) * 0.5 + float(clearance) |
| overlap_x = gap[..., 0] - abs_delta[..., 0] |
| overlap_y = gap[..., 1] - abs_delta[..., 1] |
| colliding = (overlap_x > 0.0) & (overlap_y > 0.0) |
| score = torch.where(colliding, overlap_x.clamp_min(0.0) + overlap_y.clamp_min(0.0), torch.full_like(overlap_x, -1.0)) |
| best_score, best_route = score.max(dim=2) |
| needs_move = best_score > 0.0 |
| if not bool(needs_move.any().detach().cpu()): |
| return |
| anchor_xy = route_xy.gather(1, best_route[:, :, None].expand(-1, -1, 2)) |
| anchor_sizes = route_sizes.gather(1, best_route[:, :, None].expand(-1, -1, 2)) |
| raw_dir = extra_xy - anchor_xy |
| norm = torch.linalg.norm(raw_dir, dim=2, keepdim=True) |
| angle = ( |
| env.env_i[:, None].float() * 1.618 |
| + extra_ids[None, :].float() * 2.399 |
| + route_ids[best_route].float() * 0.71 |
| + float(step) * 0.83 |
| ) |
| fallback = torch.stack((torch.cos(angle), torch.sin(angle)), dim=2) |
| direction = torch.where(norm > 1e-3, raw_dir / norm.clamp_min(1e-3), fallback) |
| support = (direction.abs() * anchor_sizes).sum(dim=2) * 0.5 + (direction.abs() * extra_sizes).sum(dim=2) * 0.5 |
| target = support + float(clearance) + 0.22 + 0.08 * ((extra_ids[None, :].float() + float(step)) % 3.0) |
| relocated = anchor_xy + direction * target[:, :, None] |
| env.platforms[:, route_len:, :2] = torch.where(needs_move[:, :, None], relocated, extra_xy) |
|
|
|
|
| def resolve_route_extra_grazes(env: Any, *, iterations: int = 4, clearance: float = 0.04) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| for _ in range(max(1, int(iterations))): |
| route_xy = env.platforms[:, :route_len, :2] |
| extra_xy = env.platforms[:, route_len:, :2] |
| route_sizes = env.sizes[:, :route_len] |
| extra_sizes = env.sizes[:, route_len:] |
| delta = extra_xy[:, :, None, :] - route_xy[:, None, :, :] |
| abs_delta = delta.abs() |
| gap = (extra_sizes[:, :, None, :] + route_sizes[:, None, :, :]) * 0.5 |
| overlap_x = gap[..., 0] - abs_delta[..., 0] |
| overlap_y = gap[..., 1] - abs_delta[..., 1] |
| colliding = (overlap_x > 0.0) & (overlap_y > 0.0) |
| score = torch.where(colliding, overlap_x.clamp_min(0.0) + overlap_y.clamp_min(0.0), torch.full_like(overlap_x, -1.0)) |
| best_score, best_route = score.max(dim=2) |
| needs_move = best_score > 0.0 |
| if not bool(needs_move.any().detach().cpu()): |
| return |
| chosen_delta = delta.gather(2, best_route[:, :, None, None].expand(-1, -1, 1, 2)).squeeze(2) |
| chosen_ox = overlap_x.gather(2, best_route[:, :, None]).squeeze(2).clamp_min(0.0) |
| chosen_oy = overlap_y.gather(2, best_route[:, :, None]).squeeze(2).clamp_min(0.0) |
| use_x = chosen_ox <= chosen_oy |
| sign_x = torch.where(chosen_delta[:, :, 0] >= 0.0, 1.0, -1.0) |
| sign_y = torch.where(chosen_delta[:, :, 1] >= 0.0, 1.0, -1.0) |
| push = torch.zeros_like(extra_xy) |
| push[:, :, 0] = torch.where(use_x, sign_x * (chosen_ox + float(clearance)), torch.zeros_like(chosen_ox)) |
| push[:, :, 1] = torch.where(~use_x, sign_y * (chosen_oy + float(clearance)), torch.zeros_like(chosen_oy)) |
| env.platforms[:, route_len:, :2] = extra_xy + push * needs_move[:, :, None].float() |
|
|
|
|
| def scatter_colliding_extras(env: Any, *, iterations: int = 6, clearance: float = 0.28) -> None: |
| route_len, p = int(env.route_len), int(env.p) |
| extras = p - route_len |
| if extras <= 0: |
| return |
| extra_ids = torch.arange(extras, device=env.device) |
| for step in range(max(1, int(iterations))): |
| route_xy = env.platforms[:, :route_len, :2] |
| extra_xy = env.platforms[:, route_len:, :2] |
| route_sizes = env.sizes[:, :route_len] |
| extra_sizes = env.sizes[:, route_len:] |
|
|
| route_delta = extra_xy[:, :, None, :] - route_xy[:, None, :, :] |
| route_abs = route_delta.abs() |
| route_gap = (extra_sizes[:, :, None, :] + route_sizes[:, None, :, :]) * 0.5 + float(clearance) |
| route_overlap_x = route_gap[..., 0] - route_abs[..., 0] |
| route_overlap_y = route_gap[..., 1] - route_abs[..., 1] |
| route_colliding = (route_overlap_x > 0.0) & (route_overlap_y > 0.0) |
| route_score = torch.where( |
| route_colliding, |
| route_overlap_x.clamp_min(0.0) + route_overlap_y.clamp_min(0.0), |
| torch.full_like(route_overlap_x, -1.0), |
| ) |
| best_route_score, best_route = route_score.max(dim=2) |
|
|
| extra_delta = extra_xy[:, :, None, :] - extra_xy[:, None, :, :] |
| extra_abs = extra_delta.abs() |
| extra_gap = (extra_sizes[:, :, None, :] + extra_sizes[:, None, :, :]) * 0.5 + float(clearance) |
| extra_overlap_x = extra_gap[..., 0] - extra_abs[..., 0] |
| extra_overlap_y = extra_gap[..., 1] - extra_abs[..., 1] |
| eye = torch.eye(extras, dtype=torch.bool, device=env.device)[None, :, :] |
| extra_colliding = (extra_overlap_x > 0.0) & (extra_overlap_y > 0.0) & ~eye |
| needs_move = (best_route_score > 0.0) | extra_colliding.any(dim=2) |
| if not bool(needs_move.any().detach().cpu()): |
| return |
|
|
| nearest_route = torch.linalg.norm(route_delta, dim=3).argmin(dim=2) |
| anchor_idx = torch.where(best_route_score > 0.0, best_route, nearest_route) |
| anchor_xy = route_xy.gather(1, anchor_idx[:, :, None].expand(-1, -1, 2)) |
| anchor_sizes = route_sizes.gather(1, anchor_idx[:, :, None].expand(-1, -1, 2)) |
| angle = ( |
| env.env_i[:, None].float() * 1.61803398875 |
| + extra_ids[None, :].float() * 2.39996322973 |
| + anchor_idx.float() * 0.73 |
| + float(step) * 0.91 |
| ) |
| direction = torch.stack((torch.cos(angle), torch.sin(angle)), dim=2) |
| support = (direction.abs() * anchor_sizes).sum(dim=2) * 0.5 + (direction.abs() * extra_sizes).sum(dim=2) * 0.5 |
| lane = ((extra_ids[None, :].float() + float(step)) % 5.0) * 0.18 |
| target = support + float(clearance) + 0.92 + lane |
| relocated = anchor_xy + direction * target[:, :, None] |
| env.platforms[:, route_len:, :2] = torch.where(needs_move[:, :, None], relocated, extra_xy) |
|
|
|
|
| def mapgen_summary(env: Any) -> dict[str, float | int]: |
| route = env.platforms[:, : env.route_len] |
| route_sizes = env.sizes[:, : env.route_len] |
| route_delta_xy = route[:, 1:, :2] - route[:, :-1, :2] |
| step_xy = torch.linalg.norm(route_delta_xy, dim=2) |
| direction = route_delta_xy / step_xy[:, :, None].clamp_min(1e-4) |
| edge_support = ( |
| (direction.abs() * route_sizes[:, :-1]).sum(dim=2) * 0.5 |
| + (direction.abs() * route_sizes[:, 1:]).sum(dim=2) * 0.5 |
| ) |
| route_edge_gap = step_xy - edge_support |
| step_z = route[:, 1:, 2] - route[:, :-1, 2] |
| route_path_xy = step_xy.sum(dim=1) |
| height_span = env.platforms[:, :, 2].amax(dim=1) - env.platforms[:, :, 2].amin(dim=1) |
| goal_xy = torch.linalg.norm(route[:, -1, :2] - route[:, 0, :2], dim=1) |
| goal_xy_to_route_path = goal_xy / route_path_xy.clamp_min(1e-4) |
| goal_z_delta = route[:, -1, 2] - route[:, 0, 2] |
| turn_abs = torch.zeros((), device=env.device) |
| if int(env.route_len) > 2: |
| prev_dir = direction[:, :-1] |
| next_dir = direction[:, 1:] |
| dot = (prev_dir * next_dir).sum(dim=2).clamp(-1.0, 1.0) |
| cross = prev_dir[..., 0] * next_dir[..., 1] - prev_dir[..., 1] * next_dir[..., 0] |
| turn_abs = torch.atan2(cross.abs(), dot).mean() |
| goal_vec = route[:, -1, :2] - route[:, 0, :2] |
| goal_dir = goal_vec / torch.linalg.norm(goal_vec, dim=1, keepdim=True).clamp_min(1e-4) |
| goal_dot = (direction * goal_dir[:, None, :]).sum(dim=2) |
| away_goal_fraction = (goal_dot < -0.10).float().mean() |
| side_goal_fraction = (goal_dot.abs() < 0.25).float().mean() |
| down_step_fraction = (step_z < -0.22).float().mean() |
| margin = route_reach_margin(env) |
| extra_reachable = torch.ones((), device=env.device) |
| extra_to_route_reachable = torch.ones((), device=env.device) |
| route_extra_predecessor_fraction = torch.zeros((), device=env.device) |
| late_route_extra_predecessor_fraction = torch.zeros((), device=env.device) |
| route_extra_predecessor_mean = torch.zeros((), device=env.device) |
| low_extra_fraction = torch.zeros((), device=env.device) |
| extra_before_start_fraction = torch.zeros((), device=env.device) |
| extra_after_goal_fraction = torch.zeros((), device=env.device) |
| extra_between_start_goal_fraction = torch.zeros((), device=env.device) |
| extra_axis_projection_mean = torch.zeros((), device=env.device) |
| lever_per_map = torch.zeros((), device=env.device) |
| locked_per_map = torch.zeros((), device=env.device) |
| locked_initially_open = torch.ones((), device=env.device) |
| if hasattr(env, "lever_mask") and hasattr(env, "locked_by"): |
| lever_per_map = env.lever_mask.float().sum(dim=1).mean() |
| locked = env.locked_by >= 0 |
| locked_per_map = locked.float().sum(dim=1).mean() |
| if bool(locked.any().detach().cpu()): |
| locked_initially_open = env.platform_available[locked].float().mean() |
| if int(env.p) > int(env.route_len): |
| route_pos = env.platforms[:, : env.route_len] |
| route_size = env.sizes[:, : env.route_len] |
| extra_pos = env.platforms[:, env.route_len :] |
| extra_size = env.sizes[:, env.route_len :] |
| route_floor = route_pos[:, :, 2].amin(dim=1, keepdim=True) |
| low_extra_fraction = (extra_pos[:, :, 2] < route_floor - 0.85).float().mean() |
| axis_vec = route_pos[:, -1:, :2] - route_pos[:, 0:1, :2] |
| axis_len = torch.linalg.norm(axis_vec, dim=2, keepdim=True).clamp_min(1e-4) |
| axis = axis_vec / axis_len |
| projection = ((extra_pos[:, :, :2] - route_pos[:, 0:1, :2]) * axis).sum(dim=2, keepdim=True) / axis_len |
| extra_axis_projection_mean = projection.mean() |
| extra_before_start_fraction = (projection < 0.0).float().mean() |
| extra_after_goal_fraction = (projection > 1.0).float().mean() |
| extra_between_start_goal_fraction = ((projection >= 0.0) & (projection <= 1.0)).float().mean() |
| route_to_extra = pair_reach_margin( |
| env, |
| route_pos[:, :, None, :], |
| extra_pos[:, None, :, :], |
| route_size[:, :, None, :], |
| extra_size[:, None, :, :], |
| ) |
| extra_to_route = pair_reach_margin( |
| env, |
| extra_pos[:, :, None, :], |
| route_pos[:, None, :, :], |
| extra_size[:, :, None, :], |
| route_size[:, None, :, :], |
| ) |
| extra_reachable = (route_to_extra.amax(dim=1) > 0.0).float().mean() |
| extra_to_route_reachable = (extra_to_route.amax(dim=2) > 0.0).float().mean() |
| extra_pred = extra_to_route > 0.0 |
| pred_counts = extra_pred.float().sum(dim=1) |
| route_nodes = torch.arange(int(env.route_len), device=env.device)[None, :] |
| inner_route = (route_nodes > 0) & (route_nodes < int(env.route_len) - 1) |
| late_route = route_nodes >= max(1, int(env.route_len) - 4) |
| route_extra_predecessor_fraction = ((pred_counts > 0.0) & inner_route).float().mean() |
| late_route_extra_predecessor_fraction = ((pred_counts > 0.0) & late_route).float().mean() |
| route_extra_predecessor_mean = pred_counts[:, 1:-1].mean() |
| footprint_area = env.sizes[:, :, 0] * env.sizes[:, :, 1] |
| return { |
| "envs": int(env.n), |
| "platforms": int(env.p), |
| "overlaps": overlap_count(env.platforms, env.sizes, margin=0.0), |
| "overlaps_2cm": overlap_count(env.platforms, env.sizes, margin=0.02), |
| "route_overlaps": route_overlap_count(env, margin=0.0), |
| "extra_overlaps": extra_overlap_count(env, margin=0.0), |
| "route_extra_overlaps": route_extra_overlap_count(env, margin=0.0), |
| "extra_from_route_reachable_fraction": float(extra_reachable.detach().cpu()), |
| "extra_to_route_reachable_fraction": float(extra_to_route_reachable.detach().cpu()), |
| "route_extra_predecessor_fraction": float(route_extra_predecessor_fraction.detach().cpu()), |
| "late_route_extra_predecessor_fraction": float(late_route_extra_predecessor_fraction.detach().cpu()), |
| "route_extra_predecessor_mean": float(route_extra_predecessor_mean.detach().cpu()), |
| "low_extra_fraction": float(low_extra_fraction.detach().cpu()), |
| "lever_per_map": float(lever_per_map.detach().cpu()), |
| "locked_per_map": float(locked_per_map.detach().cpu()), |
| "locked_initially_open_fraction": float(locked_initially_open.detach().cpu()), |
| "route_min_reach_margin": float(margin.min().detach().cpu()), |
| "route_mean_reach_margin": float(margin.mean().detach().cpu()), |
| "route_edge_gap_mean": float(route_edge_gap.mean().detach().cpu()), |
| "route_edge_gap_min": float(route_edge_gap.min().detach().cpu()), |
| "route_edge_gap_max": float(route_edge_gap.max().detach().cpu()), |
| "step_xy_mean": float(step_xy.mean().detach().cpu()), |
| "step_xy_min": float(step_xy.min().detach().cpu()), |
| "step_xy_max": float(step_xy.max().detach().cpu()), |
| "route_path_xy_mean": float(route_path_xy.mean().detach().cpu()), |
| "goal_xy_to_route_path_ratio_mean": float(goal_xy_to_route_path.mean().detach().cpu()), |
| "step_z_min": float(step_z.min().detach().cpu()), |
| "step_z_max": float(step_z.max().detach().cpu()), |
| "platform_area_mean": float(footprint_area.mean().detach().cpu()), |
| "platform_area_min": float(footprint_area.min().detach().cpu()), |
| "platform_area_max": float(footprint_area.max().detach().cpu()), |
| "tiny_platform_fraction": float((footprint_area < 0.55).float().mean().detach().cpu()), |
| "height_span_mean": float(height_span.mean().detach().cpu()), |
| "height_span_max": float(height_span.max().detach().cpu()), |
| "goal_xy_mean": float(goal_xy.mean().detach().cpu()), |
| "goal_xy_min": float(goal_xy.min().detach().cpu()), |
| "extra_axis_projection_mean": float(extra_axis_projection_mean.detach().cpu()), |
| "extra_before_start_fraction": float(extra_before_start_fraction.detach().cpu()), |
| "extra_after_goal_fraction": float(extra_after_goal_fraction.detach().cpu()), |
| "extra_between_start_goal_fraction": float(extra_between_start_goal_fraction.detach().cpu()), |
| "route_turn_abs_mean": float(turn_abs.detach().cpu()), |
| "route_away_goal_fraction": float(away_goal_fraction.detach().cpu()), |
| "route_side_goal_fraction": float(side_goal_fraction.detach().cpu()), |
| "route_down_step_fraction": float(down_step_fraction.detach().cpu()), |
| "goal_z_delta_mean": float(goal_z_delta.mean().detach().cpu()), |
| "goal_z_delta_abs_mean": float(goal_z_delta.abs().mean().detach().cpu()), |
| "goal_z_delta_min": float(goal_z_delta.min().detach().cpu()), |
| "goal_z_delta_max": float(goal_z_delta.max().detach().cpu()), |
| **all_platform_path_diagnostics(env), |
| } |
|
|
