Buckets:
| """cwm3/models/branches.py — the three base branches of CWM3 (plan §3), each trained alone first. | |
| VehicleBranch V : TCN(6 -> 256) over the CAN history [B,L,6] = (steer, gasPressed, brakePressed, vEgo, aEgo standardised, Δt) ⊕ mode | |
| embedding -> flat MLP head -> [B,25,35] (3 MDNs + gas/brake/steer_phase/resp chains). The TCN is the wm3 | |
| `tcn(cin, 256)` verbatim (cwm2/models/encoders.TCN: causal left padding, k=5, dilations 1,2,4,8). | |
| L = 40 for both variants: can4s = X (t-3.9..t @10 Hz), can8s = concat(Xpre, X)[:, 1::2] (t-7.8..t @5 Hz); | |
| Δt carries the real step offsets so the same module sees both grids. | |
| DriverBranch U : the kin_predictability GRU recipe — GRU(18 -> 128) over 40 steps of [Kh standardised 11, Kh valid bit, X standardised 5, | |
| Δt] ⊕ mode ⊕ Linear(LayerNorm(cabin PCA-256) -> 64) · cabin_ok -> flat MLP head (kin recipe: Linear-GELU-Linear) | |
| -> [B,25,26]: residual kinematics head (ΔK mu, log sig; K̂ = K_0 + ΔK̂) + hands_on / head_turned chains. | |
| cabin = none | vc256 (full-frame causal) | crop256 (route-median NON-causal, diagnostic). | |
| in_ln (2026-09-11, recipe key `u_in_ln`, default False): LayerNorm over the 12 kinematics channels (11 standardised Kh + the valid bit) | |
| of every history step at the trunk entry — the CAN channels and Δt pass through — so the trunk no longer depends on the | |
| kinematics input scale (stats v1 -> v2 changed it 3x). Off -> bitwise the old module (no extra parameters). | |
| EnvBranch E : road token bank over time -> 2-layer causal Transformer (d 128, 4 heads, FFN 256) with a learned readout query | |
| ⊕ Linear(LayerNorm(Vf PCA-256) -> 128) · vf_ok ⊕ mode -> 6 demand chains [B,25,12] + F head [B,5,256]. | |
| tokens = region (4 PCA-64 tokens per 0.5 s slot, tok4s = 8 slots / tok8s = 16 slots) | pooled (1 PCA-256 token per slot, | |
| 16 slots) | spat256 (the 4 region PCA-64 tokens of a slot FLATTENED into one 256-d token, 16 slots — keeps the spatial | |
| layout that pooling before PCA loses; 2026-09-10). Missing slots are zero tokens + a learned "missing" flag (never a key-padding mask -> no all-masked rows); | |
| causality across slots is a boolean attention mask (slot j attends to slots <= j; the readout attends to all). | |
| p_head (2026-09-11, variants `pooled_P` / `regionpca256_P`): a second flat head on the same z -> the six road-pressure chains | |
| P (ttc, ttc_drop, cut_in, lead_braking, lane_excursion, badas) [B,25,12], trunk-supervised (lambda P); off -> the old module. | |
| d_side (2026-09-11): width of the per-anchor side vector `vf` — 256 = vf256 (default) | 512 = the `regionpca256` road stream | |
| (PCA-256 of the flattened 16x1408 region grid, last window (+) 8 s mean) fed through the SAME LN -> Linear(d_side -> 128) -> | |
| GELU path in place of vf256; the token sequence / Transformer are those of `pooled` (nothing tiled or split). | |
| mode ∈ {0 manual, 1 engaged, 2 unknown} is an input embedding in every branch (one model for both regimes). | |
| can_gate (2026-09-14, U, recipe `u_can_gate`, plan Part A §A.5 arm (b), decisions.md 2026-09-15T01:48:39Z): CanReleaseGate — one scalar gate g = σ(a·act + b) | |
| (a = b = 0 at init -> 2g = 1 exactly) from the CAN activity act = mean |one-step Δ| of the standardised steer / gasPressed / brakePressed columns of Kx | |
| over the last 1 s (steps 30..39, all <= t); emitted as "HG" = 2g [B,1] and applied to the four driver-chain hazards by heads.hazards_gated / | |
| driver_loss_gated (λ' = clamp(λ·2g)). Created LAST, no RNG consumed; off (default) -> the old module bit-for-bit. | |
| adv_devices (2026-09-11, U, recipe `u_adv`): gradient-reversal device head on the GRU state (models/identity.py); f_state (E, recipe `f_state_sup`): | |
| Linear(256 -> 10) state read-out on F_hat(t+k). Both default off -> the old modules bit-for-bit. | |
| No openpilot DMS field is read anywhere in this file (inputs: CAN whitelist, YOLO-pose kinematics, V-JEPA2 latents). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| from . import heads as hd | |
| from .identity import grad_reverse # opt-in identity arm (u_adv); no-op unless adv_devices > 0 | |
| H = hd.H | |
| N_MODE = 3 | |
| class TCN(nn.Module): | |
| """wm3 train_wm3.build_net.tcn (identical to cwm2/models/encoders.TCN): Conv1d proj + 4 residual dilated blocks, last step.""" | |
| def __init__(self, cin, dim, dropout=0.1): | |
| super().__init__() | |
| self.proj = nn.Conv1d(cin, dim, 1) | |
| self.blocks = nn.ModuleList([nn.Sequential( | |
| nn.Conv1d(dim, dim, 5, padding=4 * dd, dilation=dd), nn.GELU(), nn.Dropout(dropout)) for dd in (1, 2, 4, 8)]) | |
| def forward(self, x): # x [B,L,C] | |
| h = self.proj(x.transpose(1, 2)) | |
| for blk in self.blocks: | |
| h = h + blk(h)[:, :, :h.shape[2]] # causal (left-padded) part only | |
| return h[:, :, -1] | |
| def flat_head(d_in, hidden, d_out, dropout=0.0): | |
| layers = [nn.Linear(d_in, hidden), nn.GELU()] | |
| if dropout > 0: | |
| layers.append(nn.Dropout(dropout)) | |
| layers.append(nn.Linear(hidden, d_out)) | |
| return nn.Sequential(*layers) | |
| def vproj(d_in, d_out, dropout=0.1): | |
| return nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, d_out), nn.GELU(), nn.Dropout(dropout)) | |
| class VehicleBranch(nn.Module): | |
| name = "V" | |
| def __init__(self, d=256, d_mode=16, hidden=256, dropout=0.1, cin=6): | |
| super().__init__() | |
| self.tcn = TCN(cin, d, dropout) | |
| self.mode = nn.Embedding(N_MODE, d_mode) | |
| self.head = flat_head(d + d_mode, hidden, H * hd.D_V, dropout) | |
| def forward(self, b): | |
| z = self.tcn(b["X"].float()) | |
| z = torch.cat([z, self.mode(b["mode"].long())], 1) | |
| return {"V": self.head(z).view(-1, H, hd.D_V)} | |
| class CanReleaseGate(nn.Module): | |
| """CAN release gate (2026-09-14, plan Part A §A.5 arm (b), recipe `u_can_gate`; pre-registration docs/decisions.md 2026-09-15T01:48:39Z) — the owner's | |
| heuristic ("fuse the CAN control signals with the keypoints; the driver is mostly inactive") as two parameters on the driver-chain hazards. | |
| act_t = (1/9) · Σ_{k=31..39} Σ_{c ∈ {steer, gasPressed, brakePressed}} |Kx[k, 12+c] − Kx[k−1, 12+c]| | |
| = the mean absolute one-step (0.1 s) difference of the standardised, clamped CAN control channels already in Kx (train/batch.prep_U columns 12..14) | |
| over the last 1 s of the history (steps 30..39, all <= t: causal by construction — no future CAN enters Kx); | |
| g_t = σ(a · act_t + b), scalar a = b = 0 at init -> g = 0.5 -> forward returns 2g = 1 exactly (identity on the hazards; tests/test_can_gate.py). | |
| Output "HG" = 2g [B, 1]; models/heads.hazards_gated / driver_loss_gated apply λ' = clamp(λ · 2g, 1e-5, 1 − 1e-5) to the four driver-chain hazards | |
| (hands_on on/off, head_turned on/off) at all 25 horizons. No RNG is consumed at construction (the parent's paired seed init is preserved).""" | |
| CAN0, CTRL, WIN = 12, (0, 1, 2), 10 # Kx layout: [0:11] kin, 11 valid, 12:17 CAN (steer, gasPressed, brakePressed, vEgo, aEgo), 17 Δt; WIN = 1 s @ 10 Hz | |
| def __init__(self): | |
| super().__init__() | |
| self.a = nn.Parameter(torch.zeros(1)) | |
| self.b = nn.Parameter(torch.zeros(1)) | |
| def activity(cls, kx): | |
| c = kx.float()[:, -cls.WIN:, [cls.CAN0 + j for j in cls.CTRL]] # [B, 10, 3]: the last 1 s of the three control channels | |
| return c.diff(dim=1).abs().sum(-1).mean(-1, keepdim=True) # [B, 1] | |
| def forward(self, kx): | |
| return 2.0 * torch.sigmoid(self.a * self.activity(kx) + self.b) # 2g [B, 1]; a = b = 0 -> exactly 1.0 | |
| class DriverBranch(nn.Module): | |
| name = "U" | |
| N_KIN_IN = 12 # 11 standardised Kh channels + the valid bit (train/batch.prep_U layout) | |
| def __init__(self, cin=18, hid=128, d_mode=16, cabin="none", d_cab=64, hidden=128, dropout=0.0, in_ln=False, adv_devices=0, n_face=0, gaze_chain=False, | |
| can_gate=False, n_kin=11): | |
| super().__init__() | |
| # n_kin (2026-09-18, kin v2p arm; pre-reg docs/decisions.md 2026-09-18T05:34:40Z; default 11 = bitwise the old module): width of the standardised Kh block of | |
| # train/batch.prep_U -> N_KIN_IN = n_kin + 1 (kin + valid bit, the input-LN span) and the GRU input = n_kin + 7 (+ valid, 5 CAN, Δt). Every older config has no `n_kin`. | |
| self.n_kin = int(n_kin) | |
| if self.n_kin != 11: | |
| assert not can_gate, "CanReleaseGate assumes the 18-wide Kx layout (CAN0 = 12)" | |
| self.N_KIN_IN = self.n_kin + 1 | |
| cin = self.n_kin + 7 | |
| self.in_ln = nn.LayerNorm(self.N_KIN_IN) if in_ln else None | |
| # n_face (2026-09-12, recipe `u_face`, default 0 = bitwise the old module): the causal face-landmark channels (F standardised + 1 valid bit) are appended | |
| # to every history step by train/batch.prep_U -> the GRU input widens to cin + F + 1; the LN path still covers the first 12 kinematics channels only. | |
| self.n_face = int(n_face) | |
| self.gru = nn.GRU(cin + (self.n_face + 1 if self.n_face else 0), hid, batch_first=True) | |
| self.mode = nn.Embedding(N_MODE, d_mode) | |
| self.cabin = cabin | |
| self.cab = vproj(256, d_cab, 0.1) if cabin != "none" else None | |
| d_in = hid + d_mode + (d_cab if cabin != "none" else 0) | |
| self.head = flat_head(d_in, hidden, H * hd.D_U, dropout) # kin_predictability head form: Linear-GELU-Linear | |
| # identity arm (2026-09-11, recipe `u_adv`, models/identity.py): gradient-reversal device head on the GRU state; created LAST so that every other | |
| # module keeps the paired seed init of the baseline; 0 (default) -> bitwise the old module. | |
| self.adv = nn.Linear(hid, int(adv_devices)) if adv_devices else None | |
| # gaze chain (2026-09-12, recipe `u_gaze_chain`, models/heads.G_CHAINS): a separate flat head on the same z -> the gaze_off 2-state chain [B,25,2] | |
| # (h_on, h_off logits); created LAST so the paired seed init of gru / mode / head is that of the baseline; False (default) -> bitwise the old module. | |
| self.gaze_head = flat_head(d_in, hidden, H * 2, dropout) if gaze_chain else None | |
| # CAN release gate (2026-09-14, recipe `u_can_gate`, plan Part A §A.5 arm (b)): two zero parameters on the driver-chain hazards (CanReleaseGate above); created | |
| # LAST and RNG-free so every other module keeps the paired seed init of the baseline; False (default) -> bitwise the old module. | |
| self.can_gate = CanReleaseGate() if can_gate else None | |
| def forward(self, b): | |
| kx = b["Kx"].float() | |
| if self.in_ln is not None: | |
| kx = torch.cat([self.in_ln(kx[..., :self.N_KIN_IN]), kx[..., self.N_KIN_IN:]], -1) | |
| _, h = self.gru(kx) | |
| parts = [h[-1], self.mode(b["mode"].long())] | |
| if self.cab is not None: | |
| parts.append(self.cab(b["cab"].float()) * b["cab_ok"].float()[:, None]) | |
| z = torch.cat(parts, 1) | |
| out = {"U": self.head(z).view(-1, H, hd.D_U)} | |
| if self.adv is not None: | |
| out["ADV"] = self.adv(grad_reverse(h[-1])) # device logits; the trunk sees the negated gradient (scale = lambda adv in the loss) | |
| if self.gaze_head is not None: | |
| out["G"] = self.gaze_head(z).view(-1, H, 2) | |
| if self.can_gate is not None: | |
| out["HG"] = self.can_gate(b["Kx"]) # 2g [B,1] from the last 1 s of the CAN control columns of Kx (pre-LN; the LN never touches them) | |
| return out | |
| class EnvBranch(nn.Module): | |
| name = "E" | |
| def __init__(self, tokens="region", n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, d_mode=16, d_vf=128, hidden=256, dropout=0.1, d_side=256, | |
| p_head=False, f_state=0): | |
| super().__init__() | |
| self.tokens, self.n_slots, self.d_side = tokens, n_slots, int(d_side) | |
| self.n_reg = 4 if tokens == "region" else 1 | |
| d_tok = 64 if tokens == "region" else 256 | |
| 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(self.n_reg, d) | |
| self.missing = nn.Parameter(torch.zeros(d)) | |
| self.dt_in = nn.Linear(1, d) | |
| self.readout = nn.Parameter(torch.randn(d) * 0.02) | |
| layer = nn.TransformerEncoderLayer(d, n_heads, d_ff, dropout, activation="gelu", batch_first=True, norm_first=True) | |
| self.enc = nn.TransformerEncoder(layer, n_layers) | |
| self.norm = nn.LayerNorm(d) | |
| self.vf = vproj(self.d_side, d_vf, dropout) # LN(d_side) -> Linear(d_side -> d_vf) -> GELU -> Dropout, gated by vf_ok | |
| self.mode = nn.Embedding(N_MODE, d_mode) | |
| d_z = d + d_vf + d_mode | |
| self.head = flat_head(d_z, hidden, H * hd.D_E, dropout) | |
| self.f_head = flat_head(d_z, hidden, 5 * hd.D_F, dropout) | |
| self.p_head = flat_head(d_z, hidden, H * hd.D_P, dropout) if p_head else None | |
| # state-supervised F read-out (2026-09-11, recipe `f_state_sup`, models/identity.py): Linear(D_F -> f_state) shared over the 5 horizons on F_hat; | |
| # created LAST (paired init preserved); 0 (default) -> bitwise the old module. | |
| self.f_state = nn.Linear(hd.D_F, int(f_state)) if f_state else None | |
| L = n_slots * self.n_reg + 1 | |
| slot_of = torch.arange(L) // self.n_reg | |
| slot_of[-1] = n_slots # readout row: attends to everything | |
| mask = slot_of[None, :] > slot_of[:, None] # True = key slot later than query slot -> blocked | |
| self.register_buffer("attn_mask", mask, persistent=False) | |
| def forward(self, b): | |
| tok = b["tok"].float() # [B,L,R,d_tok] | |
| B, L, R, _ = tok.shape | |
| assert L == self.n_slots and R == self.n_reg | |
| x = self.tok_in(tok) | |
| miss = (~b["tok_ok"].bool())[:, :, None, None].float() # [B,L,1,1] | |
| x = x * (1 - miss) + self.missing * miss | |
| x = x + self.slot_emb.weight[None, :, None, :] + self.reg_emb.weight[None, None, :, :] | |
| x = x + self.dt_in(b["tok_dt"].float()[:, :, None, None]) | |
| x = x.reshape(B, L * R, -1) | |
| x = torch.cat([x, self.readout[None, None].expand(B, 1, -1)], 1) | |
| h = self.enc(x, mask=self.attn_mask) | |
| z_e = self.norm(h[:, -1]) | |
| z_vf = self.vf(b["vf"].float()) * b["vf_ok"].float()[:, None] | |
| z = torch.cat([z_e, z_vf, self.mode(b["mode"].long())], 1) | |
| out = {"E": self.head(z).view(-1, H, hd.D_E), "F": self.f_head(z).view(-1, 5, hd.D_F)} | |
| if self.p_head is not None: | |
| out["P"] = self.p_head(z).view(-1, H, hd.D_P) | |
| if self.f_state is not None: | |
| out["FS"] = self.f_state(out["F"]) # [B,5,f_state] state logits at t+1..t+5 s read from F_hat | |
| return out | |
| def build_branch(branch, seed, **kw): | |
| """torch.manual_seed(seed) immediately before construction (paired init across variants of one seed).""" | |
| torch.manual_seed(seed) | |
| arch = kw.pop("arch", None) # 2026-09-12 (GPU-baselines B): architecture-swap variants live in models/branches_alt.py; absent / None -> the reference modules below, unchanged | |
| if arch: | |
| from .branches_alt import build_alt | |
| return build_alt(branch, arch, **kw) | |
| if branch == "V": | |
| return VehicleBranch(**kw) | |
| if branch == "U": | |
| return DriverBranch(**kw) | |
| if branch == "E": | |
| return EnvBranch(**kw) | |
| raise ValueError(branch) | |
| def params_by_module(model): | |
| out = {name: sum(p.numel() for p in mod.parameters() if p.requires_grad) for name, mod in model.named_children()} | |
| out.update({name: p.numel() for name, p in model.named_parameters(recurse=False)}) | |
| out["total"] = sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| return out | |
Xet Storage Details
- Size:
- 16.3 kB
- Xet hash:
- b13189973ee454a77f037e67fac3e1dfa8e6c50bb8a22b945da4ba2fd9cd1635
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.