| 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() |
|
|
