Buckets:
| """cwm3/models/interaction.py — directed, ZERO-INITIALISED residual interactions between the three base branches (plan §3, I). | |
| Every module returns a residual that is EXACTLY zero at initialisation (its last Linear has zero weight and zero bias), so the joint | |
| model's emissions equal the independent branches' bitwise at init (tests/test_zero_init_identity.py). The residuals are added to the | |
| branch TRUNKS (before the pretrained emission heads), never to the emissions themselves: | |
| Road2Driver : the driver trunk h_u [B,128] is the query over the road token sequence h_tok [B,L,128] of the E encoder (Driver-WM | |
| external -> internal cross-attention form, one head group, LN on query / keys); attention output -> LN -> Linear -> | |
| GELU -> zero-Linear -> residual on h_u. | |
| Env2Vehicle : [z_e ; z_vf] [B,256] (E readout + front PCA projection) -> LN -> Linear -> GELU -> zero-Linear -> residual on z_v [B,256]. | |
| Driver2Vehicle: h_u -> LN -> Linear -> GELU -> zero-Linear, multiplied by a per-mode sigmoid gate (Embedding(mode) init 0 -> gate 0.5) | |
| -> residual on z_v. The gate lets the manual / engaged regimes weight the driver -> vehicle path differently. | |
| NaiveFusion : the CONTROL — concat of all trunks [z_v ; h_u ; z_e ; z_vf] (640) -> Linear -> GELU -> Linear with the default | |
| (non-zero) init, split back into per-branch residuals. It perturbs every branch at init and lets every branch read | |
| every other (cwm2's fused encoder); it is what "controlled sharing" is compared against. | |
| Parameter counts (i_hidden 640, measured by models/cwm3.py): Road2Driver 231,424 · Env2Vehicle 329,088 · Driver2Vehicle 247,680 | |
| -> I = 808,192 (plan ≈ 0.8 M); NaiveFusion 821,760. No openpilot DMS field is read anywhere (inputs are the branch trunks). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| N_MODE = 3 | |
| def zero_linear(d_in, d_out): | |
| lin = nn.Linear(d_in, d_out) | |
| nn.init.zeros_(lin.weight) | |
| nn.init.zeros_(lin.bias) | |
| return lin | |
| class ResidualMLP(nn.Module): | |
| """LN -> Linear -> GELU -> Dropout -> zero-Linear. Output is exactly 0 at init for every finite input.""" | |
| def __init__(self, d_in, hidden, d_out, dropout=0.1): | |
| super().__init__() | |
| self.norm = nn.LayerNorm(d_in) | |
| self.fc1 = nn.Linear(d_in, hidden) | |
| self.drop = nn.Dropout(dropout) | |
| self.fc2 = zero_linear(hidden, d_out) | |
| def forward(self, x): | |
| return self.fc2(self.drop(F.gelu(self.fc1(self.norm(x))))) | |
| class Road2Driver(nn.Module): | |
| name = "road2drv" | |
| def __init__(self, d_u=128, d_tok=128, d_attn=128, n_heads=4, hidden=640, dropout=0.1): | |
| super().__init__() | |
| assert d_attn % n_heads == 0 | |
| self.n_heads, self.d_attn = n_heads, d_attn | |
| self.q_norm = nn.LayerNorm(d_u) | |
| self.kv_norm = nn.LayerNorm(d_tok) | |
| self.q = nn.Linear(d_u, d_attn) | |
| self.k = nn.Linear(d_tok, d_attn) | |
| self.v = nn.Linear(d_tok, d_attn) | |
| self.out = nn.Linear(d_attn, d_attn) | |
| self.mlp = ResidualMLP(d_attn, hidden, d_u, dropout) | |
| self.last_attn = None # [B, n_heads, L] of the last forward (analysis only) | |
| def forward(self, h_u, h_tok): | |
| B, L, _ = h_tok.shape | |
| nh, dh = self.n_heads, self.d_attn // self.n_heads | |
| q = self.q(self.q_norm(h_u)).view(B, 1, nh, dh).transpose(1, 2) # [B,nh,1,dh] | |
| kv = self.kv_norm(h_tok) | |
| k = self.k(kv).view(B, L, nh, dh).transpose(1, 2) # [B,nh,L,dh] | |
| v = self.v(kv).view(B, L, nh, dh).transpose(1, 2) | |
| att = torch.softmax((q @ k.transpose(-1, -2)) / dh ** 0.5, -1) # [B,nh,1,L] | |
| self.last_attn = att[:, :, 0].detach() | |
| a = (att @ v).transpose(1, 2).reshape(B, self.d_attn) | |
| return self.mlp(self.out(a)) | |
| class Env2Vehicle(nn.Module): | |
| name = "env2veh" | |
| def __init__(self, d_e=256, d_v=256, hidden=640, dropout=0.1): | |
| super().__init__() | |
| self.mlp = ResidualMLP(d_e, hidden, d_v, dropout) | |
| def forward(self, z_e_cat): | |
| return self.mlp(z_e_cat) | |
| class Driver2Vehicle(nn.Module): | |
| name = "drv2veh" | |
| def __init__(self, d_u=128, d_v=256, hidden=640, dropout=0.1, n_mode=N_MODE): | |
| super().__init__() | |
| self.mlp = ResidualMLP(d_u, hidden, d_v, dropout) | |
| self.gate = nn.Embedding(n_mode, d_v) | |
| nn.init.zeros_(self.gate.weight) # sigmoid(0) = 0.5 for every mode at init | |
| def forward(self, h_u, mode): | |
| return torch.sigmoid(self.gate(mode.long())) * self.mlp(h_u) | |
| class NaiveFusion(nn.Module): | |
| """Control: all trunks concatenated through a shared 2-layer MLP with the DEFAULT (non-zero) init; residuals to every trunk.""" | |
| name = "naive" | |
| def __init__(self, dims=(256, 128, 256), hidden=640, dropout=0.1): | |
| super().__init__() | |
| self.dims = tuple(dims) | |
| d = sum(dims) | |
| self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, d)) | |
| def forward(self, z_v, h_u, z_e_cat): | |
| r = self.net(torch.cat([z_v, h_u, z_e_cat], 1)) | |
| return torch.split(r, self.dims, 1) | |
| class RegionTokenEmbed(nn.Module): | |
| """Raw region token bank [B,L,R,d_tok] (+ tok_ok, tok_dt) -> key / value sequence [B,L*R,d] for the road -> driver cross-attention when | |
| the E base is the pooled variant: LN -> Linear(d_tok -> d) + slot / region embeddings + Δt + a learned missing vector (the EnvBranch input | |
| embedding form, freshly initialised; 11,008 params for 64 -> 128, 16 slots, 4 regions).""" | |
| def __init__(self, d_tok=64, d=128, n_slots=16, n_reg=4): | |
| super().__init__() | |
| self.n_slots, self.n_reg = n_slots, n_reg | |
| self.tok_in = nn.Sequential(nn.LayerNorm(d_tok), nn.Linear(d_tok, d)) | |
| self.slot_emb = nn.Embedding(n_slots, d) | |
| self.reg_emb = nn.Embedding(n_reg, d) | |
| self.dt_in = nn.Linear(1, d) | |
| self.missing = nn.Parameter(torch.zeros(d)) | |
| def forward(self, b): | |
| tok = b["tok"].float() | |
| B, L, R, _ = tok.shape | |
| assert L == self.n_slots and R == self.n_reg, (tok.shape, self.n_slots, self.n_reg) | |
| x = self.tok_in(tok) | |
| miss = (~b["tok_ok"].bool())[:, :, None, None].float() | |
| x = x * (1 - miss) + self.missing * miss | |
| x = x + self.slot_emb.weight[None, :, None, :] + self.reg_emb.weight[None, None, :, :] + self.dt_in(b["tok_dt"].float()[:, :, None, None]) | |
| return x.reshape(B, L * R, -1) | |
| def n_params(m): | |
| return sum(p.numel() for p in m.parameters() if p.requires_grad) | |
Xet Storage Details
- Size:
- 6.71 kB
- Xet hash:
- 9fd28660bb12eceb20e7743cb1e3a25dbf9dd77cbaef47aed99022bc19df1024
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.