from __future__ import annotations from dataclasses import dataclass import math import torch import torch.nn.functional as F LOG2PI = math.log(2.0 * math.pi) @dataclass(slots=True) class ControllerOutput: mean: torch.Tensor button_logits: torch.Tensor value: torch.Tensor class BeaconController(torch.nn.Module): action_size = 5 def __init__( self, observation_size: int, *, hidden: int = 128, depth: int = 2, control_log_std_init: float = -1.25, control_log_std_min: float = -2.30, control_log_std_max: float = -0.70, button_bias: float | None = None, forward_bias: float = 1.20, jump_bias: float = -1.25, sprint_bias: float = 1.40, ): super().__init__() layers: list[torch.nn.Module] = [] for index in range(max(1, int(depth))): layers += [ torch.nn.Linear(int(observation_size) if index == 0 else int(hidden), int(hidden)), torch.nn.SiLU(), torch.nn.LayerNorm(int(hidden)), ] self.encoder = torch.nn.Sequential(*layers) self.control_head = torch.nn.Linear(int(hidden), 3) self.button_head = torch.nn.Linear(int(hidden), 2) self.value_head = torch.nn.Linear(int(hidden), 1) self.log_std = torch.nn.Parameter(torch.full((3,), float(control_log_std_init))) self.log_std_min = float(control_log_std_min) self.log_std_max = float(control_log_std_max) torch.nn.init.normal_(self.control_head.weight, mean=0.0, std=0.01) torch.nn.init.zeros_(self.control_head.bias) with torch.no_grad(): self.control_head.bias[0] = float(forward_bias) torch.nn.init.normal_(self.button_head.weight, mean=0.0, std=0.01) torch.nn.init.zeros_(self.button_head.bias) if button_bias is not None: jump_bias = float(button_bias) sprint_bias = float(button_bias) with torch.no_grad(): self.button_head.bias[0] = float(jump_bias) self.button_head.bias[1] = float(sprint_bias) torch.nn.init.zeros_(self.value_head.bias) def forward(self, observation: torch.Tensor) -> ControllerOutput: h = self.encoder(observation) return ControllerOutput( mean=torch.tanh(self.control_head(h)), button_logits=self.button_head(h).clamp(-12.0, 12.0), value=self.value_head(h).squeeze(1), ) def choose_action( self, observation: torch.Tensor, *, deterministic_controls: bool, deterministic_buttons: bool, ) -> tuple[torch.Tensor, ControllerOutput]: out = self(observation) if deterministic_controls: controls = out.mean else: controls = (out.mean + torch.randn_like(out.mean) * self.policy_log_std().exp()).clamp(-1.0, 1.0) if deterministic_buttons: buttons = (out.button_logits > 0.0).to(dtype=observation.dtype) else: buttons = torch.bernoulli(torch.sigmoid(out.button_logits)) buttons[:, 0] = buttons[:, 0] * (observation[:, 3] > 0.5).to(dtype=observation.dtype) return torch.cat((controls, buttons), dim=1), out def act(self, observation: torch.Tensor, *, deterministic: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: action, out = self.choose_action( observation, deterministic_controls=deterministic, deterministic_buttons=deterministic, ) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return action, logprob, entropy, out.value def evaluate_actions(self, observation: torch.Tensor, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: out = self(observation) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return logprob, entropy, out.value def action_stats(self, out: ControllerOutput, action: torch.Tensor, *, jump_valid: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: controls = action[:, :3].clamp(-1.0, 1.0) buttons = action[:, 3:].clamp(0.0, 1.0) if jump_valid is None: jump_valid = torch.ones_like(buttons[:, 0]) else: jump_valid = jump_valid.reshape(-1).to(dtype=buttons.dtype) log_std = self.policy_log_std() inv_std = torch.exp(-log_std) gaussian_logprob = -0.5 * (((controls - out.mean) * inv_std).square() + 2.0 * log_std + LOG2PI).sum(1) jump_logprob = -F.binary_cross_entropy_with_logits(out.button_logits[:, 0], buttons[:, 0], reduction="none") * jump_valid sprint_logprob = -F.binary_cross_entropy_with_logits(out.button_logits[:, 1], buttons[:, 1], reduction="none") gaussian_entropy = (0.5 + 0.5 * LOG2PI + log_std).sum().expand_as(gaussian_logprob) button_probs = torch.sigmoid(out.button_logits) jump_entropy = F.binary_cross_entropy_with_logits(out.button_logits[:, 0], button_probs[:, 0], reduction="none") * jump_valid sprint_entropy = F.binary_cross_entropy_with_logits(out.button_logits[:, 1], button_probs[:, 1], reduction="none") return gaussian_logprob + jump_logprob + sprint_logprob, gaussian_entropy + jump_entropy + sprint_entropy def policy_log_std(self) -> torch.Tensor: return self.log_std.clamp(self.log_std_min, self.log_std_max) class TokenAttentionController(torch.nn.Module): """Policy encoder that compares visible platform tokens before acting. Observation layout is the v37 top-k layout: base[12] + K * token[11]. It is still purely egocentric. It does not use route ids, graph edges, or platform indices; token order can be distance-sorted by the env. """ action_size = 5 base_features = 12 token_features = 11 def __init__( self, observation_size: int, *, hidden: int = 256, depth: int = 2, heads: int = 4, control_log_std_init: float = -1.25, control_log_std_min: float = -2.30, control_log_std_max: float = -0.70, forward_bias: float = 1.20, jump_bias: float = -1.25, sprint_bias: float = 1.40, ): super().__init__() self.observation_size = int(observation_size) self.hidden = int(hidden) token_cols = max(0, self.observation_size - self.base_features) if token_cols % self.token_features != 0: raise ValueError( f"TokenAttentionController requires base+K*token obs, got obs_size={self.observation_size}" ) self.topk = token_cols // self.token_features self.base_encoder = torch.nn.Sequential( torch.nn.Linear(self.base_features, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), torch.nn.Linear(self.hidden, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), ) self.token_encoder = torch.nn.Sequential( torch.nn.Linear(self.token_features, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), torch.nn.Linear(self.hidden, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), ) self.token_condition = torch.nn.Linear(self.hidden, self.hidden) self.token_refine = torch.nn.Sequential( torch.nn.LayerNorm(self.hidden), torch.nn.Linear(self.hidden, self.hidden * 2), torch.nn.GELU(), torch.nn.Linear(self.hidden * 2, self.hidden), ) self.query = torch.nn.Sequential( torch.nn.Linear(self.hidden, self.hidden), torch.nn.SiLU(), torch.nn.Linear(self.hidden, self.hidden), ) self.fusion = torch.nn.Sequential( torch.nn.Linear(self.hidden * 3, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), torch.nn.Linear(self.hidden, self.hidden), torch.nn.SiLU(), torch.nn.LayerNorm(self.hidden), ) self.control_head = torch.nn.Linear(self.hidden, 3) self.button_head = torch.nn.Linear(self.hidden, 2) self.value_head = torch.nn.Linear(self.hidden, 1) self.log_std = torch.nn.Parameter(torch.full((3,), float(control_log_std_init))) self.log_std_min = float(control_log_std_min) self.log_std_max = float(control_log_std_max) torch.nn.init.normal_(self.control_head.weight, mean=0.0, std=0.01) torch.nn.init.zeros_(self.control_head.bias) with torch.no_grad(): self.control_head.bias[0] = float(forward_bias) torch.nn.init.normal_(self.button_head.weight, mean=0.0, std=0.01) torch.nn.init.zeros_(self.button_head.bias) with torch.no_grad(): self.button_head.bias[0] = float(jump_bias) self.button_head.bias[1] = float(sprint_bias) torch.nn.init.zeros_(self.value_head.bias) def encode(self, observation: torch.Tensor) -> torch.Tensor: base = observation[:, : self.base_features] tokens = observation[:, self.base_features :].reshape(observation.shape[0], self.topk, self.token_features) token_hit = tokens[:, :, 0] > 0.5 base_h = self.base_encoder(base) tok_h = self.token_encoder(tokens) if self.topk > 0: tok_h = tok_h + self.token_condition(base_h)[:, None, :] tok_h = tok_h + self.token_refine(tok_h) q = self.query(base_h)[:, None, :] logits = (tok_h * q).sum(dim=2) / math.sqrt(float(self.hidden)) logits = logits.masked_fill(~token_hit, -1e9) weights = torch.softmax(logits, dim=1) weights = torch.where(token_hit.any(dim=1, keepdim=True), weights, torch.zeros_like(weights)) attended = (tok_h * weights[:, :, None]).sum(dim=1) max_pool = tok_h.masked_fill(~token_hit[:, :, None], -1e6).amax(dim=1) max_pool = torch.where(token_hit.any(dim=1, keepdim=True), max_pool, torch.zeros_like(max_pool)) visible_count = token_hit.float().sum(dim=1, keepdim=True) / max(1, self.topk) max_goal = tokens[:, :, 7].amax(dim=1, keepdim=True) summary = max_pool summary[:, :1] = visible_count summary[:, 1:2] = max_goal else: attended = torch.zeros_like(base_h) summary = torch.zeros_like(base_h) return self.fusion(torch.cat((base_h, attended, summary), dim=1)) def forward(self, observation: torch.Tensor) -> ControllerOutput: h = self.encode(observation) return ControllerOutput( mean=torch.tanh(self.control_head(h)), button_logits=self.button_head(h).clamp(-12.0, 12.0), value=self.value_head(h).squeeze(1), ) def choose_action( self, observation: torch.Tensor, *, deterministic_controls: bool, deterministic_buttons: bool, ) -> tuple[torch.Tensor, ControllerOutput]: out = self(observation) if deterministic_controls: controls = out.mean else: controls = (out.mean + torch.randn_like(out.mean) * self.policy_log_std().exp()).clamp(-1.0, 1.0) if deterministic_buttons: buttons = (out.button_logits > 0.0).to(dtype=observation.dtype) else: buttons = torch.bernoulli(torch.sigmoid(out.button_logits)) buttons[:, 0] = buttons[:, 0] * (observation[:, 3] > 0.5).to(dtype=observation.dtype) return torch.cat((controls, buttons), dim=1), out def act(self, observation: torch.Tensor, *, deterministic: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: action, out = self.choose_action( observation, deterministic_controls=deterministic, deterministic_buttons=deterministic, ) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return action, logprob, entropy, out.value def evaluate_actions(self, observation: torch.Tensor, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: out = self(observation) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return logprob, entropy, out.value def action_stats(self, out: ControllerOutput, action: torch.Tensor, *, jump_valid: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: return BeaconController.action_stats(self, out, action, jump_valid=jump_valid) def policy_log_std(self) -> torch.Tensor: return self.log_std.clamp(self.log_std_min, self.log_std_max) class LocalBeaconPlannerController(torch.nn.Module): action_size = BeaconController.action_size def __init__( self, observation_size: int, *, hidden: int = 128, depth: int = 2, planner_hidden: int = 128, planner_depth: int = 2, planner_max_xy: float = 32.0, planner_max_z: float = 2.2, planner_blend_bias: float = 0.0, control_log_std_init: float = -1.25, control_log_std_min: float = -2.30, control_log_std_max: float = -0.70, button_bias: float | None = None, forward_bias: float = 1.20, jump_bias: float = -1.25, sprint_bias: float = 1.40, ): super().__init__() self.runner = BeaconController( observation_size, hidden=hidden, depth=depth, control_log_std_init=control_log_std_init, control_log_std_min=control_log_std_min, control_log_std_max=control_log_std_max, button_bias=button_bias, forward_bias=forward_bias, jump_bias=jump_bias, sprint_bias=sprint_bias, ) layers: list[torch.nn.Module] = [] for index in range(max(1, int(planner_depth))): layers += [ torch.nn.Linear(int(observation_size) if index == 0 else int(planner_hidden), int(planner_hidden)), torch.nn.SiLU(), torch.nn.LayerNorm(int(planner_hidden)), ] self.planner_encoder = torch.nn.Sequential(*layers) self.planner_head = torch.nn.Linear(int(planner_hidden), 4) self.planner_max_xy = float(planner_max_xy) self.planner_max_z = float(planner_max_z) self.planner_blend_bias = float(planner_blend_bias) self.reset_planner_head(self.planner_blend_bias) def reset_planner_head(self, blend_bias: float | None = None) -> None: torch.nn.init.zeros_(self.planner_head.weight) torch.nn.init.zeros_(self.planner_head.bias) self.planner_head.bias.data[3] = float(self.planner_blend_bias if blend_bias is None else blend_bias) def planned_observation(self, observation: torch.Tensor) -> torch.Tensor: raw = self.planner_head(self.planner_encoder(observation)) delta_xy = torch.tanh(raw[:, 0:2]) * (self.planner_max_xy / 24.0) delta_z = torch.tanh(raw[:, 2:3]) * (self.planner_max_z / 4.0) blend = torch.sigmoid(raw[:, 3:4]) planned = observation.clone() planned[:, 4:6] = (observation[:, 4:6] + blend * delta_xy).clamp(-5.0, 5.0) planned[:, 6:7] = (observation[:, 6:7] + blend * delta_z).clamp(-5.0, 5.0) local_x_m = planned[:, 4:5] * 24.0 local_y_m = planned[:, 5:6] * 24.0 planned[:, 7:8] = (torch.sqrt(local_x_m.square() + local_y_m.square() + 1e-6) / 24.0).clamp(0.0, 5.0) return planned def planner_summary(self, observation: torch.Tensor) -> dict[str, float]: with torch.no_grad(): raw = self.planner_head(self.planner_encoder(observation)) delta_xy_m = torch.tanh(raw[:, 0:2]) * self.planner_max_xy delta_z_m = torch.tanh(raw[:, 2:3]) * self.planner_max_z blend = torch.sigmoid(raw[:, 3:4]) return { "planner_blend": float(blend.mean().detach().cpu()), "planner_delta_x_m": float(delta_xy_m[:, 0].mean().detach().cpu()), "planner_delta_y_m": float(delta_xy_m[:, 1].mean().detach().cpu()), "planner_delta_z_m": float(delta_z_m[:, 0].mean().detach().cpu()), "planner_delta_xy_abs_m": float(delta_xy_m.abs().sum(1).mean().detach().cpu()), } def forward(self, observation: torch.Tensor) -> ControllerOutput: return self.runner(self.planned_observation(observation)) def choose_action( self, observation: torch.Tensor, *, deterministic_controls: bool, deterministic_buttons: bool, ) -> tuple[torch.Tensor, ControllerOutput]: planned = self.planned_observation(observation) out = self.runner(planned) if deterministic_controls: controls = out.mean else: controls = (out.mean + torch.randn_like(out.mean) * self.policy_log_std().exp()).clamp(-1.0, 1.0) if deterministic_buttons: buttons = (out.button_logits > 0.0).to(dtype=observation.dtype) else: buttons = torch.bernoulli(torch.sigmoid(out.button_logits)) buttons[:, 0] = buttons[:, 0] * (observation[:, 3] > 0.5).to(dtype=observation.dtype) return torch.cat((controls, buttons), dim=1), out def act(self, observation: torch.Tensor, *, deterministic: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: action, out = self.choose_action( observation, deterministic_controls=deterministic, deterministic_buttons=deterministic, ) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return action, logprob, entropy, out.value def evaluate_actions(self, observation: torch.Tensor, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: out = self(observation) logprob, entropy = self.action_stats(out, action, jump_valid=(observation[:, 3] > 0.5).to(dtype=observation.dtype)) return logprob, entropy, out.value def action_stats(self, out: ControllerOutput, action: torch.Tensor, *, jump_valid: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: return self.runner.action_stats(out, action, jump_valid=jump_valid) def policy_log_std(self) -> torch.Tensor: return self.runner.policy_log_std()