from __future__ import annotations from dataclasses import dataclass import math from typing import Any import torch from config import V2Config from runtime import resolve_device @dataclass(slots=True) class FastStep: observation: torch.Tensor reward: torch.Tensor done: torch.Tensor info: dict[str, torch.Tensor] class FastFloatingBeaconEnv: """Pure Torch floating-platform parkour for high-throughput RL. This is intentionally a videogame-physics training backend. It keeps the actor observation egocentric: velocity, grounded, final-goal beacon, yaw, time, and nearest-platform ray hits. It does not expose route ids, graph edges, or full map tensors to the policy. """ rays = 21 ray_features = 6 token_features = 11 base_features = 12 action_size = 5 platform_x = 0.98 platform_y = 0.86 radius = 0.18 dt = 0.05 gravity = 30.0 accel = 65.0 ground_damping = 0.55 air_damping = 0.91 air_control = 0.10 max_speed = 4.3 sprint_speed_mult = 1.35 sprint_accel_mult = 1.25 turn_speed = 3.2 jump_velocity = 7.2 landing_pad_scale = 0.25 landing_top_tolerance = 0.03 landing_prev_slack = 0.12 lava_z = -1.2 max_ray_range = 8.0 fan_radians = math.tau def __init__(self, cfg: V2Config, *, envs: int | None = None, seed: int | None = None): self.cfg = cfg self.device = resolve_device(cfg.device) self.n = int(envs if envs is not None else cfg.envs) self.route_len = int(cfg.route_jumps) + 1 self.p = self.route_len + int(cfg.distractors) self.max_steps = max(120, int(self.route_len * float(getattr(cfg, "max_steps_factor", 24.0)))) self.dt = float(getattr(cfg, "dt", self.dt)) self.gravity = float(getattr(cfg, "gravity", self.gravity)) self.accel = float(getattr(cfg, "accel", self.accel)) self.max_speed = float(getattr(cfg, "max_speed", self.max_speed)) self.jump_velocity = float(getattr(cfg, "jump_velocity", self.jump_velocity)) self.turn_speed = float(getattr(cfg, "turn_speed", self.turn_speed)) self.radius = float(getattr(cfg, "agent_radius", self.radius)) self.landing_pad_scale = float(getattr(cfg, "landing_pad_scale", self.landing_pad_scale)) self.landing_top_tolerance = float(getattr(cfg, "landing_top_tolerance", self.landing_top_tolerance)) self.landing_prev_slack = float(getattr(cfg, "landing_prev_slack", self.landing_prev_slack)) self.lava_z = float(getattr(cfg, "lava_z", self.lava_z)) self.sensor_token_sort = str(getattr(cfg, "sensor_token_sort", "index")).lower() self.observe_progress = bool(getattr(cfg, "observe_progress", True)) sensor_fov_degrees = float(getattr(cfg, "sensor_fov_degrees", 360.0)) self.fan_radians = math.radians(max(10.0, min(360.0, sensor_fov_degrees))) self.gen = torch.Generator(device=self.device) self.gen.manual_seed(int(seed if seed is not None else cfg.seed)) self.env_i = torch.arange(self.n, device=self.device) self.node_i = torch.arange(self.p, device=self.device) self.agent_offset = torch.tensor([0.0, 0.0, self.radius], device=self.device) self.vel_scale = torch.tensor([6.0, 6.0, 8.0], device=self.device) self.ray_angles = torch.linspace(-self.fan_radians * 0.5, self.fan_radians * 0.5, self.rays, device=self.device) self.platforms = torch.zeros((self.n, self.p, 3), device=self.device) self.base_platforms = torch.zeros((self.n, self.p, 3), device=self.device) self.platform_velocity = torch.zeros((self.n, self.p, 3), device=self.device) self.sizes = torch.zeros((self.n, self.p, 2), device=self.device) self.pos = torch.zeros((self.n, 3), device=self.device) self.prev_pos = torch.zeros_like(self.pos) self.vel = torch.zeros_like(self.pos) self.yaw = torch.zeros(self.n, device=self.device) self.grounded = torch.ones(self.n, dtype=torch.bool, device=self.device) self.current = torch.zeros(self.n, dtype=torch.long, device=self.device) self.visited = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.done = torch.zeros(self.n, dtype=torch.bool, device=self.device) self.success = torch.zeros(self.n, dtype=torch.bool, device=self.device) self.steps = torch.zeros(self.n, dtype=torch.long, device=self.device) self.stagnation_steps = torch.zeros(self.n, dtype=torch.long, device=self.device) self.platform_available = torch.ones((self.n, self.p), dtype=torch.bool, device=self.device) self.fragile_mask = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.fragile_age = torch.full((self.n, self.p), -1, dtype=torch.long, device=self.device) self.fragile_available = torch.ones((self.n, self.p), dtype=torch.bool, device=self.device) self.launch_pad_mask = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.moving_mask = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.move_axis = torch.zeros((self.n, self.p, 2), device=self.device) self.move_amplitude = torch.zeros((self.n, self.p), device=self.device) self.move_period_steps = torch.ones((self.n, self.p), device=self.device) self.move_phase = torch.zeros((self.n, self.p), device=self.device) self.lever_mask = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.locked_by = torch.full((self.n, self.p), -1, dtype=torch.long, device=self.device) self.lock_owner = torch.zeros((self.n, self.p), dtype=torch.long, device=self.device) self.lever_unlocked = torch.zeros((self.n, self.p), dtype=torch.bool, device=self.device) self.lever_progress = torch.zeros((self.n, self.p), dtype=torch.long, device=self.device) self.initial_goal_dist = torch.ones(self.n, device=self.device) self.best_goal_dist = torch.ones(self.n, device=self.device) self.last_progress = torch.zeros(self.n, device=self.device) self.reset() @property def obs_size(self) -> int: if self.uses_token_sensor: return self.base_features + self.sensor_topk * self.token_features return self.base_features + self.rays * self.ray_features @property def sensor_topk(self) -> int: return max(1, int(getattr(self.cfg, "sensor_topk", 16))) @property def uses_token_sensor(self) -> bool: return str(getattr(self.cfg, "sensor_mode", "ray")).lower() in {"token", "tokens", "topk", "nearest"} @property def uses_fragile_platforms(self) -> bool: return bool(getattr(self.cfg, "fragile_platforms", False)) @property def fragile_delay_steps(self) -> int: return max(1, int(getattr(self.cfg, "fragile_delay_steps", 14))) @property def uses_launch_pads(self) -> bool: return bool(getattr(self.cfg, "launch_pads", False)) @property def uses_moving_platforms(self) -> bool: return bool(getattr(self.cfg, "moving_platforms", False)) @property def uses_unlock_platforms(self) -> bool: return bool(getattr(self.cfg, "unlock_platforms", False)) def reset(self) -> torch.Tensor: self.generate_maps() self.setup_platform_mechanics(self.env_i) self.reset_state(self.env_i) self.update_dynamic_platforms() return self.observe() def reset_done(self) -> torch.Tensor: ids = torch.nonzero(self.done, as_tuple=False).flatten() if ids.numel() == 0: return self.observe() if int(ids.numel()) >= self.n: return self.reset() keep = torch.nonzero(~self.done, as_tuple=False).flatten() keep_platforms = self.platforms[keep].clone() keep_base_platforms = self.base_platforms[keep].clone() keep_platform_velocity = self.platform_velocity[keep].clone() keep_sizes = self.sizes[keep].clone() keep_platform_available = self.platform_available[keep].clone() keep_fragile_mask = self.fragile_mask[keep].clone() keep_fragile_age = self.fragile_age[keep].clone() keep_fragile_available = self.fragile_available[keep].clone() keep_launch_pad_mask = self.launch_pad_mask[keep].clone() keep_moving_mask = self.moving_mask[keep].clone() keep_move_axis = self.move_axis[keep].clone() keep_move_amplitude = self.move_amplitude[keep].clone() keep_move_period_steps = self.move_period_steps[keep].clone() keep_move_phase = self.move_phase[keep].clone() keep_lever_mask = self.lever_mask[keep].clone() keep_locked_by = self.locked_by[keep].clone() keep_lock_owner = self.lock_owner[keep].clone() keep_lever_unlocked = self.lever_unlocked[keep].clone() keep_lever_progress = self.lever_progress[keep].clone() self.generate_maps() self.setup_platform_mechanics(self.env_i) self.platforms[keep] = keep_platforms self.base_platforms[keep] = keep_base_platforms self.platform_velocity[keep] = keep_platform_velocity self.sizes[keep] = keep_sizes self.platform_available[keep] = keep_platform_available self.fragile_mask[keep] = keep_fragile_mask self.fragile_age[keep] = keep_fragile_age self.fragile_available[keep] = keep_fragile_available self.launch_pad_mask[keep] = keep_launch_pad_mask self.moving_mask[keep] = keep_moving_mask self.move_axis[keep] = keep_move_axis self.move_amplitude[keep] = keep_move_amplitude self.move_period_steps[keep] = keep_move_period_steps self.move_phase[keep] = keep_move_phase self.lever_mask[keep] = keep_lever_mask self.locked_by[keep] = keep_locked_by self.lock_owner[keep] = keep_lock_owner self.lever_unlocked[keep] = keep_lever_unlocked self.lever_progress[keep] = keep_lever_progress self.reset_state(ids) self.update_dynamic_platforms() return self.observe() def reset_state(self, ids: torch.Tensor) -> None: first = self.platforms[ids, min(1, self.route_len - 1)] - self.platforms[ids, 0] self.yaw[ids] = torch.atan2(first[:, 1], first[:, 0]) direction = first[:, :2] / torch.linalg.norm(first[:, :2], dim=1, keepdim=True).clamp_min(1e-6) start_fraction = float(getattr(self.cfg, "start_forward_fraction", 0.0)) start_fraction = min(max(start_fraction, 0.0), 0.85) half_size = (self.sizes[ids, 0] * 0.5 - self.radius * 1.25).clamp_min(0.0) edge_distance = torch.minimum( half_size[:, 0] / direction[:, 0].abs().clamp_min(1e-4), half_size[:, 1] / direction[:, 1].abs().clamp_min(1e-4), ).clamp_min(0.0) start_offset = direction * (edge_distance * start_fraction)[:, None] self.pos[ids] = self.platforms[ids, 0] self.pos[ids, :2] += start_offset self.pos[ids, 2] += self.radius self.prev_pos[ids] = self.pos[ids] self.vel[ids] = 0.0 self.grounded[ids] = True self.current[ids] = 0 self.visited[ids] = False self.visited[ids, 0] = True self.done[ids] = False self.success[ids] = False self.steps[ids] = 0 self.stagnation_steps[ids] = 0 self.fragile_available[ids] = True self.fragile_age[ids] = -1 self.fragile_mask[ids] = False self.lever_unlocked[ids] = False self.lever_progress[ids] = 0 if self.uses_fragile_platforms: self.fragile_mask[ids] = True self.fragile_mask[ids, 0] = False self.fragile_mask[ids, self.route_len - 1] = False goal = self.platforms[ids, self.route_len - 1] + self.agent_offset dist = torch.linalg.norm((goal - self.pos[ids])[:, :2], dim=1).clamp_min(1.0) self.initial_goal_dist[ids] = dist self.best_goal_dist[ids] = dist self.last_progress[ids] = 0.0 self.update_platform_availability() def setup_platform_mechanics(self, ids: torch.Tensor) -> None: self.base_platforms[ids] = self.platforms[ids] self.platform_velocity[ids] = 0.0 self.launch_pad_mask[ids] = False self.moving_mask[ids] = False self.move_axis[ids] = 0.0 self.move_amplitude[ids] = 0.0 self.move_period_steps[ids] = 1.0 self.move_phase[ids] = 0.0 self.lever_mask[ids] = False self.locked_by[ids] = -1 self.lock_owner[ids] = 0 self.lever_unlocked[ids] = False self.lever_progress[ids] = 0 route_len = int(self.route_len) if route_len <= 3: return rows = ids.reshape(-1) row_count = int(rows.numel()) if self.uses_launch_pads: fraction = float(getattr(self.cfg, "launch_pad_fraction", 0.24)) node = torch.arange(route_len, device=self.device)[None, :] eligible = (node > 0) & (node < route_len - 1) rhythm = ((node + rows[:, None]) % 5) == 2 selected = eligible & (rhythm | (self.rand(row_count, route_len) < fraction * 0.20)) self.launch_pad_mask[rows, :route_len] = selected primary_span = max(1, route_len - 4) primary = 2 + (rows % primary_span) self.launch_pad_mask[rows, primary] = True gap_extra = float(getattr(self.cfg, "launch_pad_gap_extra", 0.0)) if abs(gap_extra) > 1e-6: for pad in range(1, route_len - 1): mask = self.launch_pad_mask[rows, pad] active_rows = rows[mask] delta = self.platforms[active_rows, pad + 1, :2] - self.platforms[active_rows, pad, :2] direction = delta / torch.linalg.norm(delta, dim=1, keepdim=True).clamp_min(1e-4) self.platforms[active_rows, pad + 1 : route_len, :2] += direction[:, None, :] * gap_extra self.base_platforms[rows] = self.platforms[rows] if self.uses_moving_platforms: fraction = float(getattr(self.cfg, "moving_platform_fraction", 0.18)) node = torch.arange(route_len, device=self.device)[None, :] eligible = (node > 1) & (node < route_len - 1) & ~self.launch_pad_mask[rows, :route_len] rhythm = ((node + rows[:, None]) % 6) == 3 force_every = max(0, int(getattr(self.cfg, "moving_platform_force_every", 0))) force_offset = int(getattr(self.cfg, "moving_platform_force_offset", 3)) if force_every > 0: forced = ((node - force_offset) % force_every) == 0 rhythm = rhythm | forced selected = eligible & (rhythm | (self.rand(row_count, route_len) < fraction * 0.18)) self.moving_mask[rows, :route_len] = selected transport = bool(getattr(self.cfg, "moving_platform_transport", False)) if transport: target_span = max(3.2, float(getattr(self.cfg, "moving_platform_transport_span", 5.35))) endpoint_gap = max(0.55, float(getattr(self.cfg, "moving_platform_endpoint_gap", 1.05))) for moving_node in range(2, route_len - 1): mask = self.moving_mask[rows, moving_node] active_rows = rows[mask] prev_xy = self.platforms[active_rows, moving_node - 1, :2] next_xy = self.platforms[active_rows, moving_node + 1, :2] raw = next_xy - prev_xy raw_norm = torch.linalg.norm(raw, dim=1, keepdim=True) fallback = self.platforms[active_rows, moving_node, :2] - prev_xy fallback_norm = torch.linalg.norm(fallback, dim=1, keepdim=True).clamp_min(1e-4) direction = torch.where(raw_norm > 1e-4, raw / raw_norm.clamp_min(1e-4), fallback / fallback_norm) current_span = raw_norm.squeeze(1) extra = (target_span - current_span).clamp_min(0.0) self.platforms[active_rows, moving_node + 1 : route_len, :2] += direction[:, None, :] * extra[:, None, None] next_xy = self.platforms[active_rows, moving_node + 1, :2] self.platforms[active_rows, moving_node, :2] = (prev_xy + next_xy) * 0.5 self.platforms[active_rows, moving_node, 2] = ( self.platforms[active_rows, moving_node - 1, 2] + self.platforms[active_rows, moving_node + 1, 2] ) * 0.5 self.base_platforms[rows] = self.platforms[rows] prev_i = (torch.arange(route_len, device=self.device) - 1).clamp_min(0) next_i = (torch.arange(route_len, device=self.device) + 1).clamp_max(route_len - 1) tangent = self.platforms[rows[:, None], next_i, :2] - self.platforms[rows[:, None], prev_i, :2] tangent = tangent / torch.linalg.norm(tangent, dim=2, keepdim=True).clamp_min(1e-4) axis = tangent if transport else torch.stack((-tangent[..., 1], tangent[..., 0]), dim=2) self.move_axis[rows, :route_len] = axis amp = float(getattr(self.cfg, "moving_platform_amplitude", 0.72)) if transport: amp = max(amp, target_span * 0.5 - endpoint_gap) period = float(max(8, int(getattr(self.cfg, "moving_platform_period_steps", 70)))) self.move_amplitude[rows, :route_len] = amp * (0.75 + 0.35 * self.rand(row_count, route_len)) self.move_period_steps[rows, :route_len] = period * (0.85 + 0.35 * self.rand(row_count, route_len)) self.move_phase[rows, :route_len] = self.rand(row_count, route_len) * math.tau if self.uses_unlock_platforms and self.p > route_len and route_len > 5: gate_count = min(max(1, int(getattr(self.cfg, "unlock_gate_count", 1))), self.p - route_len) 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) span_nodes = torch.arange(lock_span, device=self.device)[None, :] for gate in range(gate_count): lever_i = route_len + gate anchor = 2 + torch.remainder(rows * (5 + 2 * gate) + gate * 3, anchor_span) locked = (anchor[:, None] + 1 + span_nodes).clamp_max(route_len - 2) self.lever_mask[rows, lever_i] = True self.locked_by[rows[:, None], locked] = lever_i self.lock_owner[rows[:, None], locked] = lever_i def update_dynamic_platforms(self) -> None: if not self.uses_moving_platforms: return self.platforms.copy_(self.base_platforms) self.platform_velocity.zero_() t = self.steps.float()[:, None] period = self.move_period_steps.clamp_min(1.0) theta = t * (math.tau / period) + self.move_phase offset = torch.sin(theta) * self.move_amplitude * self.moving_mask.float() speed_per_step = torch.cos(theta) * self.move_amplitude * (math.tau / period) * self.moving_mask.float() self.platforms[:, :, :2] = self.base_platforms[:, :, :2] + self.move_axis * offset[:, :, None] self.platform_velocity[:, :, :2] = self.move_axis * (speed_per_step / max(self.dt, 1e-6))[:, :, None] def update_platform_availability(self) -> None: self.platform_available.copy_(self.fragile_available) if self.uses_unlock_platforms: self.platform_available &= (self.locked_by < 0) | self.lever_unlocked.gather(1, self.lock_owner) def drop_unsupported(self, active: torch.Tensor) -> None: unsupported = active & self.grounded & ~self.platform_available[self.env_i, self.current] self.grounded &= ~unsupported fall_speed = torch.full_like(self.vel[:, 2], -0.25) self.vel[:, 2] = torch.where(unsupported, torch.minimum(self.vel[:, 2], fall_speed), self.vel[:, 2]) def rand(self, *shape: int) -> torch.Tensor: return torch.rand(shape, device=self.device, generator=self.gen) def update_unlocks(self, active: torch.Tensor) -> torch.Tensor: if not self.uses_unlock_platforms: return torch.zeros(self.n, dtype=torch.bool, device=self.device) rows = self.env_i current = self.current already_unlocked = self.lever_unlocked[rows, current] on_locked_lever = self.lever_mask[rows, current] & ~already_unlocked touching = active & self.grounded & on_locked_lever self.lever_progress[rows, current] += touching.long() hold_steps = max(1, int(getattr(self.cfg, "unlock_hold_steps", 1))) opened = touching & (self.lever_progress[rows, current] >= hold_steps) self.lever_unlocked[rows, current] |= opened return opened def generate_maps(self) -> 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, self.platform_y], device=self.device) size_jitter = 0.94 + 0.12 * self.rand(n, p, 2) self.sizes[:] = base_size * size_jitter swap_axes = self.rand(n) < 0.5 sign_x = torch.where(self.rand(n) < 0.5, 1.0, -1.0) sign_y = torch.where(self.rand(n) < 0.5, 1.0, -1.0) def to_world(longitudinal: torch.Tensor, lateral: torch.Tensor) -> torch.Tensor: x = torch.where(swap_axes, lateral * sign_x, longitudinal * sign_x) y = torch.where(swap_axes, longitudinal * sign_y, lateral * sign_y) return torch.stack((x, y), dim=1) height_offset = max(0.0, float(getattr(self.cfg, "map_height_offset", 2.75))) vertical = max(0.35, float(getattr(self.cfg, "map_vertical_scale", 1.25))) max_h = max(height_offset + 0.5, float(getattr(self.cfg, "map_max_height", 6.0))) lane_width = max(1.36, float(getattr(self.cfg, "map_braid_lane_width", 1.55))) edge_gap = max(0.44, float(getattr(self.cfg, "map_min_edge_gap", 0.66))) long_gap = max(1.55, self.platform_x + edge_gap) lane_gap = max(lane_width, self.platform_y + edge_gap * 0.85) mesh_lanes = max(4, int(getattr(self.cfg, "map_braid_lanes", 4)) * 2) lateral = torch.zeros(n, device=self.device) progress = torch.zeros(n, device=self.device) center_long = torch.zeros((n, route_len), device=self.device) center_lat = torch.zeros((n, route_len), device=self.device) route_long = torch.zeros((n, route_len), device=self.device) route_lat = torch.zeros((n, route_len), device=self.device) z = torch.full((n,), height_offset + 0.25, device=self.device) self.platforms[:, 0, 2] = z route_z = torch.zeros((n, route_len), device=self.device) route_z[:, 0] = z loop_period = max(4, int(getattr(self.cfg, "map_braid_loopback_period", 6))) loop_depth = float(getattr(self.cfg, "map_braid_loopback_depth", 0.45)) loop_lat = float(getattr(self.cfg, "map_braid_loopback_lateral", 0.45)) curve_scale = float(getattr(self.cfg, "map_braid_curve_scale", 1.0)) switch_period = max(3, int(getattr(self.cfg, "map_switchback_period", 4))) style_jitter = float(getattr(self.cfg, "map_braid_style_jitter", 0.25)) env_phase = self.rand(n) * math.tau for layer in range(1, route_len): phase = (layer - 1) % (loop_period + 3) side = torch.where(((self.env_i + layer) % 2) == 0, 1.0, -1.0) base_step = long_gap * (0.95 + 0.25 * self.rand(n)) lateral_step = (self.rand(n) - 0.5) * 0.24 * curve_scale * lane_gap if phase == loop_period - 1: lateral_step = side * (0.95 + 0.35 * self.rand(n)) * max(0.45, loop_lat) * lane_gap base_step = long_gap * (0.28 + 0.18 * self.rand(n)) elif phase == loop_period: lateral_step = side * (0.80 + 0.32 * self.rand(n)) * max(0.45, loop_lat) * lane_gap base_step = long_gap * (0.30 + 0.18 * self.rand(n)) * max(0.35, loop_depth) elif phase == loop_period + 1: lateral_step = -lateral * (0.46 + 0.18 * self.rand(n)) base_step = long_gap * (0.48 + 0.28 * self.rand(n)) elif layer % switch_period == 0: lateral_step = side * (0.85 + 0.35 * self.rand(n)) * lane_gap base_step = long_gap * (0.36 + 0.22 * self.rand(n)) prev_progress = progress progress = torch.maximum(progress + base_step, prev_progress + long_gap * 0.72) lateral = (lateral + lateral_step).clamp(-2.85 * lane_gap, 2.85 * lane_gap) dz = (self.rand(n) - 0.45) * 0.38 * vertical z = (z + dz).clamp(height_offset + 0.18, max_h) weave = torch.sin(env_phase + float(layer) * 0.93) * style_jitter center_long[:, layer] = progress center_lat[:, layer] = lateral route_long[:, layer] = progress route_lat[:, layer] = lateral + weave * 0.18 * lane_gap route_z[:, layer] = z self.platforms[:, layer, :2] = to_world(route_long[:, layer], route_lat[:, layer]) self.platforms[:, layer, 2] = z self.repair_route_spacing(layer, min_center=max(1.22, lane_gap * 0.82)) route_long[:, layer] = torch.where(swap_axes, self.platforms[:, layer, 1] * sign_y, self.platforms[:, layer, 0] * sign_x) route_lat[:, layer] = torch.where(swap_axes, self.platforms[:, layer, 0] * sign_x, self.platforms[:, layer, 1] * sign_y) center_long[:, layer] = route_long[:, layer] center_lat[:, layer] = lateral if p > route_len: extras = p - route_len inner_layers = max(1, route_len - 2) cells = max(extras, inner_layers * mesh_lanes) offset_seed = torch.randint(0, cells, (n, 1), device=self.device, generator=self.gen) cell = (torch.arange(extras, device=self.device)[None, :] + offset_seed) % cells anchors = 1 + (cell // mesh_lanes).clamp_max(inner_layers - 1) lane_slot = cell % mesh_lanes centered_slot = lane_slot.float() - (float(mesh_lanes - 1) * 0.5) lane_units = torch.where(centered_slot >= 0.0, centered_slot + 0.65, centered_slot - 0.65) anchor_long = center_long.gather(1, anchors) anchor_lat = center_lat.gather(1, anchors) anchor_z = route_z.gather(1, anchors) layer_wave = torch.sin(env_phase[:, None] + anchors.float() * 1.17 + lane_units * 0.41) stagger = (self.rand(n, extras) - 0.5) * (0.46 + 0.16 * curve_scale) * long_gap extra_long = anchor_long + stagger + ((anchors.float() % 2.0) - 0.5) * 0.16 * long_gap extra_lat = anchor_lat + lane_units * lane_gap + layer_wave * 0.14 * lane_gap self.platforms[:, route_len:, 0] = torch.where(swap_axes[:, None], extra_lat * sign_x[:, None], extra_long * sign_x[:, None]) self.platforms[:, route_len:, 1] = torch.where(swap_axes[:, None], extra_long * sign_y[:, None], extra_lat * sign_y[:, None]) decoy = self.rand(n, extras) < float(getattr(self.cfg, "map_illegal_decoy_fraction", 0.05)) legal_z = anchor_z + (self.rand(n, extras) - 0.5) * 0.34 * vertical 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)) ) self.platforms[:, route_len:, 2] = torch.where(decoy, illegal_z, legal_z).clamp(height_offset + 0.18, max_h) self.separate_xy_overlaps(iterations=int(getattr(self.cfg, "map_overlap_cleanup_iters", 24))) center = self.platforms[:, :, :2].mean(dim=1, keepdim=True) self.platforms[:, :, :2] -= center def repair_route_spacing(self, layer: int, *, min_center: float) -> None: if layer <= 1: return needs_global_spacing = ( float(getattr(self.cfg, "map_hairpin_depth", 0.0)) > 0.0 or float(getattr(self.cfg, "map_unintuitive_depth", 0.0)) > 0.0 or float(getattr(self.cfg, "map_valley_depth", 0.0)) > 0.0 or int(getattr(self.cfg, "map_redirect_period", 0)) > 0 or float(getattr(self.cfg, "map_backtrack_depth", 0.0)) > 0.0 ) lookback = self.route_len if needs_global_spacing else 5 lo = max(0, int(layer) - int(lookback)) current = self.platforms[:, layer, :2] previous = self.platforms[:, lo:layer, :2] delta = current[:, None, :] - previous dist = torch.linalg.norm(delta, dim=2).clamp_min(1e-4) too_close = dist < float(min_center) side = torch.where(((self.env_i + layer) % 2) == 0, 1.0, -1.0) fallback = torch.stack((side, -side), dim=1) direction = torch.where(dist[:, :, None] > 1e-3, delta / dist[:, :, None], fallback[:, None, :]) push = (float(min_center) - dist).clamp_min(0.0) move = (direction * push[:, :, None] * too_close[:, :, None].float()).sum(dim=1) self.platforms[:, layer, :2] = current + move.clamp(-0.65, 0.65) def separate_xy_overlaps(self, *, iterations: int = 24) -> None: if self.p <= 1: return margin = max(0.12, float(getattr(self.cfg, "map_braid_mesh_clearance", 0.16))) eye = torch.eye(self.p, dtype=torch.bool, device=self.device)[None, :, :] mobility = torch.ones((self.n, self.p), device=self.device) mobility[:, 0] = 0.08 mobility[:, self.route_len - 1] = 0.12 route_cap = min(self.route_len, self.p) mobility[:, :route_cap] *= 0.55 for _ in range(max(0, int(iterations))): xy = self.platforms[:, :, :2] delta = xy[:, :, None, :] - xy[:, None, :, :] abs_delta = delta.abs() gap_x = (self.sizes[:, :, None, 0] + self.sizes[:, None, :, 0]) * 0.5 + margin gap_y = (self.sizes[:, :, None, 1] + self.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) push = (pair_push * share[..., None]).sum(dim=2) self.platforms[:, :, :2] = xy + push.clamp(-0.48, 0.48) def basis(self) -> tuple[torch.Tensor, torch.Tensor]: forward = torch.stack((torch.cos(self.yaw), torch.sin(self.yaw)), dim=1) right = torch.stack((-torch.sin(self.yaw), torch.cos(self.yaw)), dim=1) return forward, right def observe(self) -> torch.Tensor: forward, right = self.basis() goal_pos = self.platforms[:, self.route_len - 1] + self.agent_offset goal_delta = goal_pos - self.pos goal_local = torch.stack(((goal_delta[:, :2] * forward).sum(1), (goal_delta[:, :2] * right).sum(1)), dim=1) goal_dist = torch.linalg.norm(goal_delta[:, :2], dim=1, keepdim=True) vel_local = torch.stack(((self.vel[:, :2] * forward).sum(1), (self.vel[:, :2] * right).sum(1), self.vel[:, 2]), dim=1) base = torch.cat( ( vel_local / self.vel_scale, self.grounded.float()[:, None], goal_local / 24.0, goal_delta[:, 2:3] / 4.0, goal_dist / 24.0, torch.sin(self.yaw)[:, None], torch.cos(self.yaw)[:, None], (1.0 - self.steps.float() / max(1, self.max_steps))[:, None], (self.last_progress if self.observe_progress else torch.zeros_like(self.last_progress))[:, None], ), dim=1, ) sensor = self.platform_tokens() if self.uses_token_sensor else self.sensor_fan() return torch.cat((base, sensor.reshape(self.n, -1)), dim=1).clamp(-5.0, 5.0) def platform_tokens(self) -> torch.Tensor: """Nearest local platform tokens, range-limited and egocentric. This intentionally exposes no route id, graph edge, or next-platform label. It is the same local geometry a ray sensor tries to summarize, but without collapsing multiple nearby platforms into one angular bin. """ forward, right = self.basis() surface_pos = self.platforms + self.agent_offset delta = surface_pos - self.pos[:, None] local_x = (delta[:, :, :2] * forward[:, None]).sum(2) local_y = (delta[:, :, :2] * right[:, None]).sum(2) local_z = delta[:, :, 2] dist = torch.sqrt((local_x.square() + local_y.square()).clamp_min(1e-6)) bearing = torch.atan2(local_y, local_x) in_fov = bearing.abs() <= (self.fan_radians * 0.5) max_range = max(0.5, float(getattr(self.cfg, "sensor_token_range", self.max_ray_range))) locked = self.locked_by >= 0 visible = self.platform_available | self.lever_mask | locked score = torch.where((dist <= max_range) & in_fov & visible, dist, torch.full_like(dist, 1e6)) k = min(self.sensor_topk, self.p) nearest_dist, nearest_i = torch.topk(score, k, dim=1, largest=False, sorted=True) hit = nearest_dist < 1e5 if self.sensor_token_sort in {"distance", "dist"}: order_key = nearest_dist + (~hit).float() * 1e6 elif self.sensor_token_sort in {"bearing", "angle"}: order_key = bearing.gather(1, nearest_i.clamp(0, self.p - 1)).abs() + (~hit).float() * 1e6 else: # Backwards-compatible old checkpoint layout. Do not use this for # new no-leak training because route platforms are created before # distractors, so index ordering can expose a weak route prior. order_key = nearest_i.float() + (~hit).float() * float(self.p + 1) order = torch.argsort(order_key, dim=1) nearest_dist = nearest_dist.gather(1, order) nearest_i = nearest_i.gather(1, order) hit = hit.gather(1, order) idx = nearest_i.clamp(0, self.p - 1) token_x = local_x.gather(1, idx) token_y = local_y.gather(1, idx) token_z = local_z.gather(1, idx) token_sx = self.sizes[:, :, 0].gather(1, idx) token_sy = self.sizes[:, :, 1].gather(1, idx) token_goal = (idx == (self.route_len - 1)) & hit token_available = self.platform_available.gather(1, idx) & hit token_lever = self.lever_mask.gather(1, idx) & hit token_locked = locked.gather(1, idx) & ~token_available & hit tokens = torch.stack( ( hit.float(), torch.where(hit, token_x / max_range, torch.zeros_like(token_x)), torch.where(hit, token_y / max_range, torch.zeros_like(token_y)), torch.where(hit, token_z / 3.0, torch.zeros_like(token_z)), torch.where(hit, token_sx / 1.7, torch.zeros_like(token_sx)), torch.where(hit, token_sy / 1.5, torch.zeros_like(token_sy)), torch.where(hit, nearest_dist / max_range, torch.ones_like(nearest_dist)), token_goal.float(), token_available.float(), token_lever.float(), token_locked.float(), ), dim=2, ) if k < self.sensor_topk: pad = torch.zeros((self.n, self.sensor_topk - k, self.token_features), device=self.device, dtype=tokens.dtype) tokens = torch.cat((tokens, pad), dim=1) return tokens def sensor_fan(self) -> torch.Tensor: forward, right = self.basis() surface_pos = self.platforms + self.agent_offset delta = surface_pos - self.pos[:, None] local_x = (delta[:, :, :2] * forward[:, None]).sum(2) local_y = (delta[:, :, :2] * right[:, None]).sum(2) local_z = delta[:, :, 2] dist = torch.sqrt((local_x.square() + local_y.square()).clamp_min(1e-6)) bearing = torch.atan2(local_y, local_x) diff = torch.atan2( torch.sin(bearing[:, :, None] - self.ray_angles[None, None, :]), torch.cos(bearing[:, :, None] - self.ray_angles[None, None, :]), ) half_width = torch.atan2(self.sizes.amax(dim=2) * 0.55 + self.radius, dist.clamp_min(0.25)).clamp(0.035, 0.34) locked = self.locked_by >= 0 visible = self.platform_available | self.lever_mask | locked hit_mask = (dist[:, :, None] < self.max_ray_range) & (diff.abs() <= half_width[:, :, None]) & visible[:, :, None] score = torch.where(hit_mask, dist[:, :, None], torch.full_like(diff, 1e6)) nearest_dist, nearest_i = score.min(dim=1) hit = nearest_dist < 1e5 idx = nearest_i.clamp(0, self.p - 1) hit_z = local_z.gather(1, idx) hit_size = self.sizes.amax(dim=2).gather(1, idx) hit_goal = idx == (self.route_len - 1) hit_available = self.platform_available.gather(1, idx) & hit hit_lever = self.lever_mask.gather(1, idx) & hit hit_locked = locked.gather(1, idx) & ~hit_available & hit hit_state = hit_lever.float() - hit_locked.float() return torch.stack( ( hit.float(), torch.where(hit, nearest_dist / self.max_ray_range, torch.ones_like(nearest_dist)), torch.where(hit, hit_z / 3.0, torch.zeros_like(nearest_dist)), torch.where(hit, hit_size / 1.5, torch.zeros_like(nearest_dist)), torch.where(hit, hit_goal.float(), torch.zeros_like(nearest_dist)), torch.where(hit, hit_state, torch.zeros_like(nearest_dist)), ), dim=2, ) def step(self, action: torch.Tensor) -> FastStep: self.update_dynamic_platforms() self.update_platform_availability() active = ~self.done action = action.to(self.device, dtype=torch.float32) move_forward = action[:, 0].clamp(-1.0, 1.0) move_side = action[:, 1].clamp(-1.0, 1.0) turn = action[:, 2].clamp(-1.0, 1.0) jump_signal = action[:, 3] > 0.5 sprint = action[:, 4].clamp(0.0, 1.0) self.drop_unsupported(active) carry = active & self.grounded & self.moving_mask[self.env_i, self.current] & self.platform_available[self.env_i, self.current] self.pos[:, :2] += self.platform_velocity[self.env_i, self.current, :2] * self.dt * carry[:, None].float() self.prev_pos.copy_(self.pos) self.yaw += turn * self.turn_speed * self.dt * active.float() forward, right = self.basis() move = forward * move_forward[:, None] + right * move_side[:, None] move = move / torch.linalg.norm(move, dim=1, keepdim=True).clamp_min(1.0) sprint_signal = sprint[:, None] * move_forward.clamp_min(0.0)[:, None] control = torch.where(self.grounded[:, None], torch.ones_like(move), torch.full_like(move, self.air_control)) accel = self.accel * (1.0 + (self.sprint_accel_mult - 1.0) * sprint_signal) self.vel[:, :2] += move * accel * self.dt * control * active[:, None] damping = torch.where(self.grounded[:, None], torch.full_like(self.vel[:, :2], self.ground_damping), torch.full_like(self.vel[:, :2], self.air_damping)) self.vel[:, :2] *= damping max_speed = self.max_speed * (1.0 + (self.sprint_speed_mult - 1.0) * sprint_signal) speed = torch.linalg.norm(self.vel[:, :2], dim=1, keepdim=True) self.vel[:, :2] *= torch.minimum(torch.ones_like(speed), max_speed / speed.clamp_min(1e-6)) do_jump = active & self.grounded & jump_signal self.vel[:, 2] = torch.where(do_jump, torch.full_like(self.vel[:, 2], self.jump_velocity), self.vel[:, 2]) self.grounded &= ~do_jump self.vel[:, 2] -= self.gravity * self.dt * active.float() self.pos += self.vel * self.dt * active[:, None] landed, land_i = self.landings() self.pos[:, 2] = torch.where(landed, self.platforms[self.env_i, land_i, 2] + self.radius, self.pos[:, 2]) self.vel[:, 2] = torch.where(landed, self.vel[:, 2].clamp_min(0.0), self.vel[:, 2]) self.grounded = landed self.current = torch.where(landed, land_i, self.current) first_visit = landed & ~self.visited[self.env_i, land_i] self.visited.scatter_(1, self.current[:, None], True) self.update_fragile_platforms(landed, land_i, active) unlocked_now = self.update_unlocks(active) self.apply_launch_pads(landed, land_i, active) goal = self.route_len - 1 on_goal = active & self.grounded & (self.current == goal) fallen = active & (self.pos[:, 2] < self.lava_z) self.steps += active.long() self.update_dynamic_platforms() self.update_platform_availability() timeout = active & (self.steps >= self.max_steps) previous_best = self.best_goal_dist dist = self.goal_distance() self.best_goal_dist = torch.minimum(self.best_goal_dist, dist) progress_delta = ((previous_best - self.best_goal_dist) / self.initial_goal_dist.clamp_min(1.0)).clamp_min(0.0) self.last_progress = (1.0 - self.best_goal_dist / self.initial_goal_dist.clamp_min(1.0)).clamp(0.0, 1.0) stagnation_timeout = torch.zeros_like(active) patience = int(getattr(self.cfg, "stagnation_patience_steps", 0)) if patience > 0: epsilon = float(getattr(self.cfg, "stagnation_progress_epsilon", 0.002)) warmup = int(getattr(self.cfg, "stagnation_warmup_steps", 32)) improved = (progress_delta > epsilon) | first_visit | on_goal tracked = active & (self.steps >= warmup) & ~fallen & ~on_goal self.stagnation_steps = torch.where( active & improved, torch.zeros_like(self.stagnation_steps), self.stagnation_steps + tracked.long(), ) stagnation_timeout = tracked & (self.stagnation_steps >= patience) self.success |= on_goal timeout = timeout | stagnation_timeout self.done |= on_goal | fallen | timeout self.stagnation_steps = torch.where(self.done, torch.zeros_like(self.stagnation_steps), self.stagnation_steps) reward = ( progress_delta * float(getattr(self.cfg, "goal_progress_reward_scale", 2.0)) + first_visit.float() * float(getattr(self.cfg, "first_visit_reward", 0.05)) + unlocked_now.float() * float(getattr(self.cfg, "unlock_reward", 1.5)) + on_goal.float() * 20.0 + on_goal.float() * float(getattr(self.cfg, "success_time_bonus", 0.0)) * (1.0 - self.steps.float() / max(1, self.max_steps)).clamp(0.0, 1.0) - fallen.float() * float(getattr(self.cfg, "fall_penalty", 8.0)) - timeout.float() * 2.0 - stagnation_timeout.float() * float(getattr(self.cfg, "stagnation_penalty", 0.0)) - active.float() * float(getattr(self.cfg, "step_cost", 0.002)) ) return FastStep( self.observe(), reward, self.done.clone(), { "success": on_goal, "fallen": fallen, "timeout": timeout, "stagnation": stagnation_timeout, "unlocked": unlocked_now, "progress": self.last_progress, "stage": self.last_progress * max(1, self.route_len - 1), }, ) def update_fragile_platforms(self, landed: torch.Tensor, land_i: torch.Tensor, active: torch.Tensor) -> None: if not self.uses_fragile_platforms: return rows = self.env_i touched_fragile = active & landed & self.fragile_mask[rows, land_i] new_touch = touched_fragile & (self.fragile_age[rows, land_i] < 0) self.fragile_age[rows[new_touch], land_i[new_touch]] = 0 ticking = (self.fragile_age >= 0) & self.platform_available & active[:, None] self.fragile_age = torch.where(ticking, self.fragile_age + 1, self.fragile_age) collapsed = self.fragile_mask & self.fragile_available & (self.fragile_age >= self.fragile_delay_steps) self.fragile_available &= ~collapsed self.update_platform_availability() self.drop_unsupported(active) def apply_launch_pads(self, landed: torch.Tensor, land_i: torch.Tensor, active: torch.Tensor) -> None: if not self.uses_launch_pads: return launch = active & landed & self.launch_pad_mask[self.env_i, land_i] velocity = float(getattr(self.cfg, "launch_pad_velocity", 10.4)) self.vel[:, 2] = torch.where(launch, torch.full_like(self.vel[:, 2], velocity), self.vel[:, 2]) self.grounded &= ~launch self.pos[:, 2] = torch.where(launch, self.pos[:, 2] + 0.03, self.pos[:, 2]) def landings(self) -> tuple[torch.Tensor, torch.Tensor]: foot_prev = self.prev_pos[:, 2] - self.radius foot_now = self.pos[:, 2] - self.radius top = self.platforms[:, :, 2] crossing = ( (self.vel[:, 2:3] <= 0.0) & (foot_prev[:, None] >= top - self.landing_prev_slack) & (foot_now[:, None] <= top + self.landing_top_tolerance) ) travel = (foot_prev - foot_now).clamp_min(1e-6) alpha = ((foot_prev[:, None] - top) / travel[:, None]).clamp(0.0, 1.0) hit_xy = self.prev_pos[:, None, :2] + alpha[:, :, None] * (self.pos[:, None, :2] - self.prev_pos[:, None, :2]) delta_xy = (hit_xy - self.platforms[:, :, :2]).abs() inside = (delta_xy <= self.sizes * 0.5 + self.radius * self.landing_pad_scale).all(dim=2) candidate = crossing & inside & self.platform_available score = torch.where(candidate, top, torch.full_like(top, -1e9)) land_i = score.argmax(dim=1) return candidate.any(dim=1), land_i def goal_distance(self) -> torch.Tensor: goal = self.platforms[:, self.route_len - 1] + self.agent_offset return torch.linalg.norm((goal - self.pos)[:, :2], dim=1) def progress_tensor(self) -> torch.Tensor: return self.last_progress def snapshots(self, indices: list[int] | torch.Tensor) -> list[dict[str, Any]]: if isinstance(indices, torch.Tensor): ids = [int(x) for x in indices.detach().cpu().tolist()] else: ids = [int(x) for x in indices] rows: list[dict[str, Any]] = [] platforms = self.platforms.detach().cpu() sizes = self.sizes.detach().cpu() pos = self.pos.detach().cpu() yaw = self.yaw.detach().cpu() steps = self.steps.detach().cpu() grounded = self.grounded.detach().cpu() done = self.done.detach().cpu() success = self.success.detach().cpu() fallen = (self.pos[:, 2] < self.lava_z).detach().cpu() visited = self.visited.detach().cpu() available = self.platform_available.detach().cpu() fragile = self.fragile_mask.detach().cpu() launch_pad = self.launch_pad_mask.detach().cpu() moving = self.moving_mask.detach().cpu() lever_mask = self.lever_mask.detach().cpu() lever_unlocked = self.lever_unlocked.detach().cpu() locked_by = self.locked_by.detach().cpu() lever_progress = self.lever_progress.detach().cpu() platform_velocity = self.platform_velocity.detach().cpu() stage = (self.last_progress * max(1, self.route_len - 1)).detach().cpu() for i in ids: timeout = bool(done[i] and (not success[i]) and (not fallen[i])) outcome = "success" if bool(success[i]) else ("fallen" if bool(fallen[i]) else ("timeout" if timeout else "running")) rows.append( { "pos": pos[i].tolist(), "yaw": float(yaw[i]), "stage": float(stage[i]), "step": int(steps[i]), "grounded": bool(grounded[i]), "success": bool(success[i]), "done": bool(done[i]), "fallen": bool(fallen[i]), "burned": False, "timeout": timeout, "outcome": outcome, "landed": int(self.current[i].detach().cpu()), "goal": self.route_len - 1, "egoReplay": True, "physicsBackend": "torch_fast_kinematic", "routeMode": "route_free_goal", "plannerMode": "policy_direct_ego_sensor", "sensorMode": str(getattr(self.cfg, "sensor_mode", "ray")), "platforms": platforms[i].tolist(), "sizes": sizes[i].tolist(), "activeNodes": [bool(x) for x in available[i].tolist()], "visited": [bool(x) for x in visited[i].tolist()], "fireMask": [False for _ in range(self.p)], "fireActive": [False for _ in range(self.p)], "leverMask": [bool(x) for x in lever_mask[i].tolist()], "leverUnlocked": [bool(x) for x in lever_unlocked[i].tolist()], "leverProgress": [int(x) for x in lever_progress[i].tolist()], "leverHoldSteps": int(max(1, int(getattr(self.cfg, "unlock_hold_steps", 1)))), "lockedBy": [int(x) for x in locked_by[i].tolist()], "fragileMask": [bool(x) for x in fragile[i].tolist()], "launchPadMask": [bool(x) for x in launch_pad[i].tolist()], "movingMask": [bool(x) for x in moving[i].tolist()], "platformVelocity": platform_velocity[i].tolist(), "platformOpen": [bool(x) for x in available[i].tolist()], "platformAvailable": [bool(x) for x in available[i].tolist()], "platformSafe": [bool(x) for x in available[i].tolist()], } ) return rows