| """Strict-local temporal routing variants for the PAIR observability pilot. |
| |
| Both variants keep the frozen Local-ARCA policy body. A causal local history |
| encoder can only change Local-ARCA's action-query router. PAIR-Belief differs |
| from History-only solely by a training-time cosine target from the validated |
| synchronized-action complementarity teacher. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from stereo_decoder_variants import StereoARCA |
|
|
|
|
| class StereoHistoryARCA(StereoARCA): |
| def __init__(self, *args, history=32, history_stride=4, role_dim=64, **kwargs): |
| state_dim = int(args[0] if args else kwargs["state_dim"]) |
| super().__init__(*args, **kwargs) |
| d = self.query.shape[-1] |
| self.history_length = int(history) |
| self.history_stride = int(history_stride) |
| self.history_input = nn.Sequential( |
| nn.LayerNorm(d + state_dim), |
| nn.Linear(d + state_dim, d), nn.GELU(), |
| ) |
| self.history_gru = nn.GRU(d, d, num_layers=2, batch_first=True, dropout=.1) |
| self.route_history = nn.Linear(d, d, bias=False) |
| self.belief_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, role_dim)) |
| nn.init.zeros_(self.route_history.weight) |
| self._online_history = None |
| self.last_belief = None |
|
|
| def reset_history(self): |
| self._online_history = None |
|
|
| def _encode_history(self, history): |
| encoded = self.history_input(history) |
| _, hidden = self.history_gru(encoded) |
| summary = hidden[-1] |
| self.last_belief = F.normalize(self.belief_head(summary), dim=-1) |
| return summary |
|
|
| def _online_window(self, observation, qpos): |
| current = torch.cat((observation.mean(1).detach(), qpos.detach()), -1) |
| if self._online_history is None or len(self._online_history) != len(current): |
| self._online_history = [[] for _ in range(len(current))] |
| windows = [] |
| for index, row in enumerate(current): |
| self._online_history[index].append(row) |
| maximum = 1 + (self.history_length - 1) * self.history_stride |
| self._online_history[index] = self._online_history[index][-maximum:] |
| values = self._online_history[index][::-self.history_stride][::-1] |
| pad = [torch.zeros_like(row) for _ in range(self.history_length - len(values))] |
| windows.append(torch.stack(pad + values)) |
| return torch.stack(windows) |
|
|
| def _route_with_history(self, state, observation, history_summary, batch): |
| q = self.query.expand(batch, -1, -1) |
| context = (self.route_state(state) + self.route_observation(observation.mean(1)) |
| + self.route_history(history_summary)) |
| features = self.route_mlp(q + context.unsqueeze(1)) |
| logits = torch.matmul(features, self.role_prototypes.t()) / features.shape[-1] ** .5 |
| values, ids = logits.topk(2, dim=-1) |
| gates = torch.zeros_like(logits).scatter_(-1, ids, values.softmax(-1).to(logits.dtype)) |
| importance = logits.softmax(-1).mean((0, 1)) |
| load = gates.gt(0).to(logits.dtype).mean((0, 1)) / 2.0 |
| return gates, self.roles_n * (importance * load).sum() |
|
|
| def forward(self, image, depth_mm, qpos, actions=None, history=None): |
| x = self._rgbd_tokens(image, depth_mm) |
| state_vec = self.state(qpos) |
| if history is None: |
| history = self._online_window(x, qpos) |
| summary = self._encode_history(history.to(dtype=x.dtype)) |
| gates, aux = self._route_with_history(state_vec, x, summary, image.shape[0]) |
| if actions is not None: |
| h = self.posterior(self.action(actions) + self.pos) |
| mu, logvar = self.latent(h.mean(1)).chunk(2, -1) |
| logvar = logvar.clamp(-10., 5.) |
| z = mu + torch.randn_like(mu) * torch.exp(.5 * logvar) |
| else: |
| mu = logvar = None |
| z = torch.zeros((image.shape[0], self.z_proj.in_features), device=image.device) |
| memory = torch.cat((state_vec.unsqueeze(1), self.z_proj(z).unsqueeze(1), x), dim=1) |
| decoded = self.decoder(self.query.expand(image.shape[0], -1, -1), memory, x, gates) |
| return self.out(decoded), mu, logvar, aux, self.last_belief |
|
|
