| 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: |
| |
| |
| |
| 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 |
|
|
