Buckets:
| """cwm3/models/u_driverwm.py — U trunk option `driverwm` (owner experiment C, pre-registration docs/decisions.md 2026-09-17T07:00:22Z): the Driver-WM | |
| temporal encoder of the AIDE port (aide/models/driverwm_kp.py, INTERNAL stream) as the BATON driver branch, at its PUBLISHED width — NOT parameter-matched | |
| (stated in every table that reports it) — consuming EXACTLY the inputs of the GRU U branch and producing an h_u of the GRU's width so that the flat head, the | |
| mode embedding, the joint model's decoders / K head / interactions and every evaluator see the contract they already know. | |
| Inputs (unchanged; train/batch.prep_U): Kx [B, 40, 18] = 11 standardised Kh kinematics channels + valid bit + 5 standardised CAN columns + Δt over the causal | |
| history t-3.9 .. t @ 10 Hz, mode ∈ {0, 1, 2}. The input LayerNorm over the 12 kinematics channels (recipe `u_in_ln`, DriverBranch.in_ln), the mode embedding | |
| and the flat Linear-GELU-Linear head are the reference modules, constructed FIRST under the paired seed (models/branches.build_branch), so they start from the | |
| same initial values as the GRU baseline of the same seed; only the sequence encoder differs. | |
| Architecture (== driverwm_kp's internal stream with the external stream absent — the main model `all_minus_road2drv` has no road -> driver path, so the gated | |
| injection g · m_t of the exterior context has no counterpart here; this is the paper's own do(Ext = ∅) setting): | |
| tok Linear(18 -> D_INT = 1024) per history step + learned step embedding [40, 1024] (driverwm_kp: PartTokenizer 4 x 256 concatenated = 1024 + step emb) | |
| pre_int ONE causal pre-norm nn.TransformerEncoderLayer(1024, 4 heads, d_ff 4096, dropout 0) (driverwm_kp.pre_int: same class, width, heads, d_ff, mask) | |
| z_t the encoder state at the last step (= t); every step j attends to steps <= j only (upper-triangular boolean mask, driverwm_repro.causal_mask) | |
| f_int the residual MLP transition LN(1024) -> 4096 -> GELU -> 1024 (driverwm_repro.res_mlp), applied ONCE: z~ = z_t + f_int(z_t) — the Driver-WM dynamics | |
| core reduced to a single latent step (its 5-step roll-out decodes one pose per step; our head emits all 25 horizons from one state, as the GRU does) | |
| proj LayerNorm(1024) -> Linear(1024 -> 128) = h_u (the only module without a Driver-WM counterpart; needed for the contract) | |
| Parameter count is reported by `DriverBranchDWM.arch_info` (encoder ≈ 21.1 M vs the reference GRU(18 -> 128) 56,832; the whole U branch ≈ 21.2 M vs 159,314). | |
| The joint model (models/cwm3.CWM3Joint) reads `d_trunk` and calls `encode_trunk(kx)` when the U module exposes them; every GRU base keeps its old path bit for bit. | |
| No openpilot DMS field is read anywhere (inputs: CAN whitelist columns, YOLO-pose kinematics). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| from . import heads as hd | |
| from .branches import DriverBranch | |
| from .identity import grad_reverse | |
| H = hd.H | |
| L_HIST = 40 | |
| D_INT, HEADS, D_FF = 1024, 4, 4096 # the AIDE port's published internal width: 4 parts x 256 = 1024, 4 heads, d_ff 4096 (aide/models/driverwm_kp.py) | |
| def n_params(m): | |
| return sum(p.numel() for p in m.parameters()) | |
| def causal_mask(L, device=None): | |
| """True = blocked (key later than query) — driverwm_repro.causal_mask.""" | |
| return torch.triu(torch.ones(L, L, dtype=torch.bool, device=device), 1) | |
| def res_mlp(D, hidden): | |
| """driverwm_repro.res_mlp verbatim: LN -> Linear -> GELU -> Linear (residual added by the caller).""" | |
| return nn.Sequential(nn.LayerNorm(D), nn.Linear(D, hidden), nn.GELU(), nn.Linear(hidden, D)) | |
| class DriverWMEncoder(nn.Module): | |
| """Driver-WM internal stream on the kinematics history: tokeniser + step embedding -> causal pre-norm Transformer layer -> last state -> f_int -> LN -> Linear.""" | |
| def __init__(self, cin, d_out, d=D_INT, heads=HEADS, d_ff=D_FF, L=L_HIST, dropout=0.0, transition=True): | |
| super().__init__() | |
| self.d, self.L, self.transition = int(d), int(L), bool(transition) | |
| self.tok = nn.Linear(cin, d) | |
| self.pos = nn.Parameter(torch.zeros(L, d)) | |
| nn.init.normal_(self.pos, std=0.02) # driverwm_kp: step / view embeddings ~ N(0, 0.02) | |
| self.pre_int = nn.TransformerEncoderLayer(d, heads, d_ff, dropout=dropout, batch_first=True, norm_first=True) # == driverwm_kp.pre_int (GELU is not used there either: default relu) | |
| self.f_int = res_mlp(d, d_ff) if transition else None | |
| self.out_ln = nn.LayerNorm(d) | |
| self.proj = nn.Linear(d, d_out) | |
| self.register_buffer("mask", causal_mask(L), persistent=False) | |
| def steps(self, x): # [B,L,C] -> causal per-step states [B,L,d] (state j depends on steps <= j only) | |
| L = x.shape[1] | |
| assert L <= self.L, (L, self.L) | |
| z = self.tok(x) + self.pos[None, :L] | |
| return self.pre_int(z, src_mask=self.mask[:L, :L]) | |
| def forward(self, x): # [B,L,C] -> h_u [B,d_out] | |
| z = self.steps(x)[:, -1] | |
| if self.f_int is not None: | |
| z = z + self.f_int(z) # one latent step of the Driver-WM dynamics core (g = 0: no exterior stream) | |
| return self.proj(self.out_ln(z)) | |
| class DriverBranchDWM(DriverBranch): | |
| """DriverBranch with the GRU replaced by the Driver-WM internal-stream encoder (published width). Same forward contract as DriverBranch / DriverBranchAlt.""" | |
| 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, transition=True): | |
| super().__init__(cin, hid, d_mode, cabin, d_cab, hidden, dropout, in_ln, adv_devices, n_face, gaze_chain, can_gate) # reference modules first: paired init kept | |
| cin_eff = self.gru.input_size | |
| ref = n_params(self.gru) | |
| del self.gru | |
| self.enc = DriverWMEncoder(cin_eff, hid, D_INT, HEADS, D_FF, L_HIST, 0.0, transition) | |
| self.d_trunk = int(hid) | |
| parts = {"tok": n_params(self.enc.tok) + self.enc.pos.numel(), "pre_int": n_params(self.enc.pre_int), | |
| "f_int": n_params(self.enc.f_int) if self.enc.f_int is not None else 0, "proj": n_params(self.enc.out_ln) + n_params(self.enc.proj)} | |
| self.arch_info = {"branch": "U", "arch": "driverwm", "d_int": D_INT, "heads": HEADS, "d_ff": D_FF, "layers": 1, "transition": bool(transition), | |
| "encoder_params": int(n_params(self.enc)), "encoder_parts": parts, "reference_encoder_params": int(ref), | |
| "ratio": round(n_params(self.enc) / ref, 2), "reference": "GRU(18 -> 128)", "parameter_matched": False, | |
| "source": "aide/models/driverwm_kp.py internal stream (pre_int TransformerEncoderLayer(1024, 4, 4096) + f_int res_mlp(1024, 4096)); external stream absent"} | |
| # ------------------------------------------------------------------ the trunk hooks the joint model uses (models/cwm3.CWM3Joint.trunk_U) | |
| def encode_trunk(self, kx): # kx = the (LN-applied) Kx [B,40,C] -> h_u [B,hid] | |
| return self.enc(kx) | |
| 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.encode_trunk(kx) | |
| parts = [h, 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)) | |
| 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"]) | |
| return out | |
| if __name__ == "__main__": | |
| torch.manual_seed(0) | |
| m = DriverBranchDWM(hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True) | |
| print(m.arch_info) | |
| print("branch total", n_params(m), {k: n_params(v) for k, v in m.named_children()}) | |
| b = {"Kx": torch.randn(3, L_HIST, 18), "mode": torch.tensor([0, 1, 2])} | |
| print("out", tuple(m(b)["U"].shape)) | |
Xet Storage Details
- Size:
- 8.47 kB
- Xet hash:
- a3be895035844ce2d7874b0ec21b8a5d7b706d29195a5118723997f50054db37
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.