Buckets:
| """cwm3/models/heads.py — emission layouts, likelihoods and marginals of the three base branches (plan §3). | |
| Fork of cwm2/models/heads.py (chain_bce, mdn_nll/mean/sample, gauss_nll, masked_mean are verbatim); the layouts are per branch: | |
| V (35) : [0:9] steer MDN (3 comps x (w, mu, log_sig pre-softplus)) on standardised steeringAngleDeg | |
| [9:18] vEgo MDN [18:27] aEgo MDN | |
| [27:35] chains gas, brake, steer_phase, resp x (logit h_on, logit h_off) | |
| U (26) : [0:11] ΔK mu (standardised; K̂ = K_0 + ΔK̂, K_0 = last valid history step) [11:22] K log_sig | |
| [22:26] chains hands_on, head_turned x (h_on, h_off) | |
| E (12) : chains steer_rate, decel, ttc, ttc_drop, lane, badas x (h_on, h_off); the F head is a separate [5,256] output. | |
| P (12) : road-pressure chains ttc, ttc_drop, cut_in, lead_braking, lane_excursion, badas x (h_on, h_off) — the E0 `*_P` variants' extra output | |
| (2026-09-11; same chain form as D, loss weight lambda P, part names chain_P_<name> -> part of the selection objective). | |
| Events for the pred_layout contract: on-hazard of a chain = onset event, off-hazard = offset event | |
| (hand_return = hands_on.on, hand_leave = hands_on.off, head_turn_onset = head_turned.on, head_return = head_turned.off). | |
| CAN release gate (2026-09-14, plan Part A §A.5 arm (b), opt-in recipe `u_can_gate`; models/branches.CanReleaseGate): hazards_gated / gate_hazards / chain_bce_h / | |
| driver_loss_gated score the two driver chains on λ' = clamp(λ · 2g, EPS_H, 1 − EPS_H); with 2g == 1 (a = b = 0 at init) they are bitwise hazards / driver_loss. | |
| No DMS field anywhere: the driver chains come from pose-derived kinematics targets. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| H = 25 | |
| KMIX = 3 | |
| V_CHAINS = ("gas", "brake", "steer_phase", "resp") | |
| U_CHAINS = ("hands_on", "head_turned") | |
| G_CHAINS = ("gaze_off",) # 2026-09-12 opt-in third driver chain (recipe `u_gaze_chain`): DriverBranch output "G" [B,25,2] = (h_on, h_off) logits | |
| E_CHAINS = ("steer_rate", "decel", "ttc", "ttc_drop", "lane", "badas") | |
| P_CHAINS = ("ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas") | |
| D_V, D_U, D_E, D_F, D_K = 27 + 2 * len(V_CHAINS), 22 + 2 * len(U_CHAINS), 2 * len(E_CHAINS), 256, 11 | |
| D_P = 2 * len(P_CHAINS) | |
| EPS_H = 1e-5 | |
| LOG2PI_HALF = 0.9189385332046727 | |
| # ----------------------------------------------------------------------------- slicing | |
| def split_V(V): | |
| V = V.float() | |
| lead = V.shape[:-1] | |
| return {"steer": V[..., 0:9].reshape(*lead, 3, KMIX), "vEgo": V[..., 9:18].reshape(*lead, 3, KMIX), | |
| "aEgo": V[..., 18:27].reshape(*lead, 3, KMIX), "chains": V[..., 27:D_V].reshape(*lead, len(V_CHAINS), 2)} | |
| def split_U(U): | |
| U = U.float() | |
| return {"dK_mu": U[..., 0:11], "K_logsig": U[..., 11:22], "chains": U[..., 22:D_U].reshape(*U.shape[:-1], len(U_CHAINS), 2)} | |
| def split_E(E): | |
| E = E.float() | |
| return {"chains": E.reshape(*E.shape[:-1], len(E_CHAINS), 2)} | |
| def split_P(P): | |
| P = P.float() | |
| return {"chains": P.reshape(*P.shape[:-1], len(P_CHAINS), 2)} | |
| def hazards(chains): | |
| """chains [..., C, 2] logits -> (h_on, h_off) probabilities clamped to (EPS_H, 1-EPS_H).""" | |
| p = torch.sigmoid(chains.float()).clamp(EPS_H, 1 - EPS_H) | |
| return p[..., 0], p[..., 1] | |
| def gate_hazards(h, gate2): | |
| """h [B, ...] clamped hazards, gate2 = 2·g [B, 1] (CAN release gate) -> clamp(h · 2g, EPS_H, 1 − EPS_H), 2g broadcast over the trailing dims. | |
| gate2 == 1 exactly -> h bitwise (x · 1.0 == x for finite x; the clamp is idempotent on an already clamped h).""" | |
| g = gate2.float().reshape(gate2.shape[0], *([1] * (h.dim() - 1))) | |
| return (h * g).clamp(EPS_H, 1 - EPS_H) | |
| def hazards_gated(chains, gate2): | |
| """CAN release gate (2026-09-14, plan Part A §A.5 arm (b); models/branches.CanReleaseGate): chains [B, T, C, 2] logits, gate2 = 2·g [B, 1] -> | |
| (h_on, h_off) = clamp(λ · 2g, EPS_H, 1 − EPS_H) with λ = hazards(chains) — the prediction-pack hazards of a gated run (train_cwm3.predict_split).""" | |
| h_on, h_off = hazards(chains) | |
| return gate_hazards(h_on, gate2), gate_hazards(h_off, gate2) | |
| # ----------------------------------------------------------------------------- likelihoods (verbatim cwm2 forms) | |
| def chain_bce(logit_on, logit_off, z_prev, z, m): | |
| """Two-state chain transition NLL per step; zero where m == 0.""" | |
| h_on = torch.sigmoid(logit_on.float()).clamp(EPS_H, 1 - EPS_H) | |
| h_off = torch.sigmoid(logit_off.float()).clamp(EPS_H, 1 - EPS_H) | |
| z, zp = z.float(), z_prev.float() | |
| l = -((1 - zp) * (z * h_on.log() + (1 - z) * (1 - h_on).log()) | |
| + zp * ((1 - z) * h_off.log() + z * (1 - h_off).log())) | |
| return l * m.float() | |
| def chain_bce_h(h_on, h_off, z_prev, z, m): | |
| """chain_bce on HAZARDS already in (EPS_H, 1 − EPS_H) (the gated λ' of gate_hazards): the arithmetic of chain_bce after its sigmoid / clamp, verbatim.""" | |
| h_on, h_off = h_on.float(), h_off.float() | |
| z, zp = z.float(), z_prev.float() | |
| l = -((1 - zp) * (z * h_on.log() + (1 - z) * (1 - h_on).log()) | |
| + zp * ((1 - z) * h_off.log() + z * (1 - h_off).log())) | |
| return l * m.float() | |
| def mdn_nll(params, y): | |
| params = params.float() | |
| w = torch.log_softmax(params[..., 0, :], -1) | |
| mu = params[..., 1, :] | |
| sig = F.softplus(params[..., 2, :]) + 1e-3 | |
| zsc = (y.float()[..., None] - mu) / sig | |
| lc = -0.5 * zsc ** 2 - sig.log() - LOG2PI_HALF | |
| return -torch.logsumexp(w + lc, -1) | |
| def mdn_mean(params): | |
| params = params.float() | |
| w = torch.softmax(params[..., 0, :], -1) | |
| return (w * params[..., 1, :]).sum(-1) | |
| def mdn_sample(params, S, gen=None): | |
| """params [B,T,3,K] -> samples [S,B,T].""" | |
| params = params.float() | |
| w = torch.softmax(params[..., 0, :], -1) | |
| mu = params[..., 1, :] | |
| sig = F.softplus(params[..., 2, :]) + 1e-3 | |
| B, T, K = w.shape | |
| comp = torch.multinomial(w.reshape(-1, K), S, replacement=True, generator=gen).reshape(B, T, S).permute(2, 0, 1) | |
| m = mu[None].expand(S, -1, -1, -1).gather(-1, comp[..., None])[..., 0] | |
| s = sig[None].expand(S, -1, -1, -1).gather(-1, comp[..., None])[..., 0] | |
| return m + s * torch.randn(m.shape, device=m.device, generator=gen) | |
| def gauss_nll(mu, logsig, y, m): | |
| logsig = logsig.float().clamp(-7, 5) | |
| sig = logsig.exp() | |
| nll = 0.5 * ((y.float() - mu.float()) / sig) ** 2 + logsig + LOG2PI_HALF | |
| return nll * m.float() | |
| def masked_mean(x, m): | |
| m = m.float() | |
| return (x * m).sum() / m.sum().clamp_min(1.0) | |
| def forward_marginal(h_on, h_off, z0): | |
| """pi_k = P(z_k = 1) by forward recursion; h_on/h_off [B,T], z0 [B] -> [B,T].""" | |
| pi, prev = [], z0.float() | |
| for k in range(h_on.shape[1]): | |
| prev = prev * (1 - h_off[:, k]) + (1 - prev) * h_on[:, k] | |
| pi.append(prev) | |
| return torch.stack(pi, 1) | |
| # ----------------------------------------------------------------------------- per-branch losses | |
| def vehicle_loss(V, tg, lam_hz=0.5): | |
| """V [B,T,35]; tg: Ystd[B,T,3] (steer, vEgo, aEgo standardised), M[B,T], zV_prev/zV/zVm [B,T,4] (gas, brake, phase, resp). | |
| -> total, parts (chain_* entries form the selection objective; nll_* are the MDN terms).""" | |
| v = split_V(V) | |
| M = tg["M"].float() | |
| parts, total = {}, 0.0 | |
| for name, ci in (("steer", 0), ("vEgo", 1), ("aEgo", 2)): | |
| nll = masked_mean(mdn_nll(v[name], tg["Ystd"][..., ci]), M) | |
| parts["nll_" + name] = nll | |
| total = total + nll | |
| c = v["chains"] | |
| for j, name in enumerate(V_CHAINS): | |
| l = masked_mean(chain_bce(c[..., j, 0], c[..., j, 1], tg["zV_prev"][..., j], tg["zV"][..., j], tg["zVm"][..., j]), tg["zVm"][..., j]) | |
| parts["chain_" + name] = l | |
| total = total + (lam_hz if name == "steer_phase" else 1.0) * l | |
| return total, parts | |
| def driver_loss(U, tg, lam_K=0.5): | |
| """U [B,T,26]; tg: K0std[B,11], Kstd[B,T,11], Km[B,T,11], zU_prev/zU/zUm [B,T,2]. -> total, parts.""" | |
| u = split_U(U) | |
| K_mu = tg["K0std"][:, None, :] + u["dK_mu"] | |
| lk = masked_mean(gauss_nll(K_mu, u["K_logsig"], tg["Kstd"], tg["Km"]), tg["Km"]) | |
| parts, total = {"nll_K": lk}, lam_K * lk | |
| c = u["chains"] | |
| for j, name in enumerate(U_CHAINS): | |
| l = masked_mean(chain_bce(c[..., j, 0], c[..., j, 1], tg["zU_prev"][..., j], tg["zU"][..., j], tg["zUm"][..., j]), tg["zUm"][..., j]) | |
| parts["chain_" + name] = l | |
| total = total + l | |
| return total, parts | |
| def driver_loss_gated(U, tg, gate2, lam_K=0.5): | |
| """driver_loss under the CAN release gate (2026-09-14, plan Part A §A.5 arm (b)): each driver chain is scored on λ' = clamp(λ · 2g) where λ is the clamped | |
| sigmoid of the chain logits computed exactly as inside chain_bce (same op order -> bitwise driver_loss while gate2 == 1); the K term and the part names are | |
| those of driver_loss (chain_hands_on / chain_head_turned enter the selection objective unchanged).""" | |
| u = split_U(U) | |
| K_mu = tg["K0std"][:, None, :] + u["dK_mu"] | |
| lk = masked_mean(gauss_nll(K_mu, u["K_logsig"], tg["Kstd"], tg["Km"]), tg["Km"]) | |
| parts, total = {"nll_K": lk}, lam_K * lk | |
| c = u["chains"] | |
| for j, name in enumerate(U_CHAINS): | |
| h_on = gate_hazards(torch.sigmoid(c[..., j, 0].float()).clamp(EPS_H, 1 - EPS_H), gate2) | |
| h_off = gate_hazards(torch.sigmoid(c[..., j, 1].float()).clamp(EPS_H, 1 - EPS_H), gate2) | |
| l = masked_mean(chain_bce_h(h_on, h_off, tg["zU_prev"][..., j], tg["zU"][..., j], tg["zUm"][..., j]), tg["zUm"][..., j]) | |
| parts["chain_" + name] = l | |
| total = total + l | |
| return total, parts | |
| def gaze_loss(G, tg): | |
| """G [B,T,2] (h_on, h_off) logits of the opt-in gaze_off chain (2026-09-12, recipe `u_gaze_chain`); tg: zG_prev / zG / zGm [B,T] from the causal gaze | |
| targets (train/batch.prep_U). -> total, parts — the part is named chain_gaze_off so that it enters selection_objective like every other chain.""" | |
| G = G.float() | |
| l = masked_mean(chain_bce(G[..., 0], G[..., 1], tg["zG_prev"], tg["zG"], tg["zGm"]), tg["zGm"]) | |
| return l, {"chain_" + G_CHAINS[0]: l} | |
| def env_loss(E, Fh, tg, lam_F=0.1, P=None, lam_P=1.0): | |
| """E [B,T,12], Fh [B,5,256]; tg: zD_prev/zD/zDm [B,T,6], F[B,5,256], Fok[B,5]. -> total, parts (L_F included in selection unless select_LF false). | |
| P [B,T,12] (optional, E0 `*_P` variants): road-pressure chains against tg zP_prev/zP/zPm, weight lam_P, parts chain_P_<name> (selection).""" | |
| e = split_E(E) | |
| c = e["chains"] | |
| parts, total = {}, 0.0 | |
| for j, name in enumerate(E_CHAINS): | |
| l = masked_mean(chain_bce(c[..., j, 0], c[..., j, 1], tg["zD_prev"][..., j], tg["zD"][..., j], tg["zDm"][..., j]), tg["zDm"][..., j]) | |
| parts["chain_" + name] = l | |
| total = total + l | |
| if P is not None and "zP" in tg: | |
| cp = split_P(P)["chains"] | |
| for j, name in enumerate(P_CHAINS): | |
| l = masked_mean(chain_bce(cp[..., j, 0], cp[..., j, 1], tg["zP_prev"][..., j], tg["zP"][..., j], tg["zPm"][..., j]), tg["zPm"][..., j]) | |
| parts["chain_P_" + name] = l | |
| total = total + lam_P * l | |
| cos = F.cosine_similarity(Fh.float(), tg["F"].float(), dim=-1) | |
| lf = lam_F * masked_mean(1 - cos, tg["Fok"]) | |
| parts["L_F"] = lf | |
| return total + lf, parts | |
| def selection_objective(parts, select_LF=True): | |
| """Validation model-selection objective: sum of Bernoulli chain NLLs (+ the bounded L_F when present and select_LF); MDN / Gaussian-K | |
| terms are excluded (cwm2 amendment 2026-09-09: unbounded below, they overfit within a few epochs and froze the chains). | |
| select_LF=False (recipe key, 2026-09-11 P2 verdict: F cosine is mostly the present's similarity to the future): L_F stays an auxiliary | |
| training loss but no longer enters model selection; old runs (select_LF absent -> True) stay reproducible.""" | |
| return float(sum(float(v) for k, v in parts.items() if k.startswith("chain_")) + (float(parts.get("L_F", 0.0)) if select_LF else 0.0)) | |
Xet Storage Details
- Size:
- 11.9 kB
- Xet hash:
- a402a886ef09505e7aa205fa04628804a1a7cc735b1b4d16f889d4ba149f3f74
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.