File size: 4,291 Bytes
20b0922 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """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
|