Buckets:
| """cwm3/models/latent_aux.py — C4 (R2 WP4 batch 2, 2026-09-17; plan <local-notes>/plans/elegant-dazzling-stallman.md 'C4 面向任务的多视界潜变量辅助目标'; | |
| pre-registration docs/decisions.md 2026-09-16T06:46:35Z + the C4 lane row of 2026-09-17): a task-oriented LOW-DIMENSIONAL multi-horizon latent auxiliary | |
| objective on the joint model's trunk features. | |
| What is predicted. From the (post-interaction) trunk features of the anchor t, x_t = [z_v ; h_u ; z_e ; z_vf] (640-d, the same vector the NaiveFusion control | |
| reads), the predictor emits one d_z-dimensional latent per horizon k in HORIZONS (5 / 10 / 15 steps = 1 / 2 / 3 s). The regression TARGET at horizon k is the | |
| EMA target encoder's embedding g_ema(s_{t+k}) of a low-dimensional, PHYSICALLY / STATE-supervised future state vector | |
| s_{t+k} in R^13 = [ 5 continuous road quantities (standardised with TRAIN statistics; results/targets_road_dyn Yf, whitelist CAN only): | |
| lead_dRel, lead_vRel, lead_vLead, lane_center_offset at t+k, and ttc_inv = clip(max(-vRel, 0) / max(dRel, 1 m), 0, 1) (1 / TTC, 0 when | |
| not approaching) — the TTC margin, | |
| 6 road-pressure occupancy bits P_{t+k} (results/targets_pressure families ttc, ttc_drop, cut_in, lead_braking, lane_excursion, badas), | |
| 2 driver-state occupancy bits (hands_on, head_turned at t+k; the v2 causal kin targets S of train/batch.prep_U) ] | |
| with a validity mask v_{t+k} in {0,1}^13 (lead channels valid iff the lead is valid at t+k; lane iff the lane is valid; P iff Pvalid; driver iff the state | |
| is observed; everything AND the s12 in-range mask M). The encoder input is [s ⊙ v ; v] (26-d), so missing components are explicit. | |
| Encoders / heads (LatentAux): | |
| pred LN(640) -> Linear(640, d_hidden) -> GELU -> Linear(d_hidden, K·d_z) the online predictor on the trunk (trainable; gradients reach the trunks) | |
| enc Linear(26, d_enc) -> GELU -> Linear(d_enc, d_z) the ONLINE target encoder g_theta (trainable by the reconstruction + variance terms only) | |
| dec Linear(d_z, 13) linear read-back of s from g_theta(s): the physical grounding of the latent | |
| enc_ema the same encoder as BUFFERS (state_dict, not parameters): g_xi <- m g_xi + (1 - m) g_theta after every optimiser step (ema_update); its output is the | |
| regression target and is computed under no_grad (stop-gradient on the TARGET side only) | |
| Loss (loss()): L_aux = L_pred + rec_coef · L_rec + var_coef · L_var | |
| L_pred masked MSE between pred_k(x_t) and g_xi(s_{t+k}) over rows with >= 1 valid component at k (updates pred AND the trunks / interaction modules) | |
| L_rec masked reconstruction of s from g_theta(s): MSE on the 5 continuous channels, BCE-with-logits on the 8 bits, normalised by the number of valid components | |
| L_var VICReg variance term on g_theta(s) over the valid rows of the batch: mean_j relu(1 - sqrt(Var_j + 1e-4)) (anti-collapse; NOT stacked with SimNorm) | |
| The trainer adds lam · L_aux to the joint loss (parts lat_pred / lat_rec / lat_var; none starts with chain_ -> the selection objective is unchanged). lam = 0 is a | |
| TRUE no-op: train_cwm3_joint.latent_enabled returns False and the module is not built at all (tests/test_latent_aux.py). The stop-gradient does NOT stop the | |
| auxiliary loss from changing the trunks: L_pred back-propagates through x_t into V / U / E and the interaction modules — this is the mechanism by which C4 can | |
| move the onset endpoints, in either direction (stated in the pre-registration). No claim of any NextLat theoretical guarantee is made. | |
| Strict future: build_targets reads Yf / Yvalid ONLY at the horizon indices (the trainer slices them to [N, K, 4] on the host), P / S at the same indices; | |
| perturbing every other step leaves s / v bitwise unchanged (test). ~0.18 M trainable parameters at the defaults (n_params()). No openpilot DMS field. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| HORIZONS = (5, 10, 15) # steps of 0.2 s -> 1 / 2 / 3 s | |
| CONT = ("lead_dRel", "lead_vRel", "lead_vLead", "lane_center_offset", "ttc_inv") | |
| P_BITS = ("ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas") | |
| U_BITS = ("hands_on", "head_turned") | |
| N_CONT, N_BITS = len(CONT), len(P_BITS) + len(U_BITS) | |
| D_S = N_CONT + N_BITS # 13 | |
| S_CLIP = 5.0 | |
| DEFAULTS = dict(d_z=16, d_hidden=256, d_enc=64, ema=0.99, var_coef=1.0, rec_coef=1.0) | |
| KW = tuple(DEFAULTS) + ("horizons",) | |
| def ttc_inv(dRel, vRel): | |
| """1 / TTC in 1/s from the lead relative distance (m) and relative speed (m/s, negative = approaching), clipped to [0, 1] (TTC >= 1 s), 0 when not approaching.""" | |
| return torch.clamp(torch.relu(-vRel) / torch.clamp(dRel, min=1.0), 0.0, 1.0) | |
| def continuous_channels(Yf): | |
| """Yf [..., 4] (lead_dRel, lead_vRel, lead_vLead, lane_center_offset) -> [..., 5] with ttc_inv appended (raw units).""" | |
| return torch.cat([Yf[..., :4], ttc_inv(Yf[..., 0], Yf[..., 1])[..., None]], -1) | |
| def continuous_valid(Yv): | |
| """Yvalid [..., 4] -> [..., 5]: ttc_inv is valid iff the lead is valid.""" | |
| return torch.cat([Yv[..., :4], Yv[..., 0:1]], -1) | |
| def build_targets(LYf, LYv, P, Pv, S, M, lat_mu, lat_sd, horizons=HORIZONS): | |
| """-> {"s": [B,K,13] f32, "v": [B,K,13] bool, "row_ok": [B,K] bool}. | |
| LYf [B,K,4] f32 / LYv [B,K,4] bool: future road quantities ALREADY sliced to the horizons (host side, train_cwm3_joint.attach_latent_targets); | |
| P / Pv [B,25,6]; S [B,25,2] (NaN = unobserved); M [B,25] in-range mask; lat_mu / lat_sd [5] TRAIN statistics of the continuous channels.""" | |
| idx = torch.as_tensor([k - 1 for k in horizons], device=LYf.device, dtype=torch.long) | |
| c = continuous_channels(LYf.float()) | |
| cv = continuous_valid(LYv.bool()) | |
| mu, sd = torch.as_tensor(lat_mu, device=LYf.device, dtype=torch.float32), torch.as_tensor(lat_sd, device=LYf.device, dtype=torch.float32) | |
| cs = torch.clamp(torch.nan_to_num((c - mu) / sd.clamp_min(1e-6)), -S_CLIP, S_CLIP) * cv.float() | |
| p = P[:, idx].float() | |
| pv = Pv[:, idx].bool() | |
| Sk = S[:, idx].float() | |
| uv = torch.isfinite(Sk) | |
| u = torch.nan_to_num(Sk) | |
| s = torch.cat([cs, p * pv.float(), u * uv.float()], -1) | |
| v = torch.cat([cv, pv, uv], -1) & M[:, idx].bool()[:, :, None] | |
| s = s * v.float() | |
| return {"s": s, "v": v, "row_ok": v.any(-1)} | |
| def latent_stats(LYf, LYv): | |
| """TRAIN mean / sd (float64) of the 5 continuous channels over valid entries; numpy in, numpy out (the trainer stores them as lat_mu / lat_sd).""" | |
| import numpy as np | |
| c = continuous_channels(torch.from_numpy(np.asarray(LYf, np.float32))).numpy().astype(np.float64) | |
| v = continuous_valid(torch.from_numpy(np.asarray(LYv, bool))).numpy() | |
| mu = np.array([c[..., j][v[..., j]].mean() if v[..., j].any() else 0.0 for j in range(N_CONT)]) | |
| sd = np.array([c[..., j][v[..., j]].std() if v[..., j].sum() > 1 else 1.0 for j in range(N_CONT)]) | |
| return {"lat_mu": mu.astype(np.float32), "lat_sd": np.maximum(sd, 1e-3).astype(np.float32)} | |
| class LatentAux(nn.Module): | |
| def __init__(self, d_in, d_z=16, d_hidden=256, d_enc=64, ema=0.99, var_coef=1.0, rec_coef=1.0, horizons=HORIZONS): | |
| super().__init__() | |
| self.horizons = tuple(int(k) for k in horizons) | |
| self.K, self.d_z, self.m = len(self.horizons), int(d_z), float(ema) | |
| self.var_coef, self.rec_coef = float(var_coef), float(rec_coef) | |
| self.pred = nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, d_hidden), nn.GELU(), nn.Linear(d_hidden, self.K * self.d_z)) | |
| self.enc = nn.Sequential(nn.Linear(2 * D_S, d_enc), nn.GELU(), nn.Linear(d_enc, self.d_z)) | |
| self.dec = nn.Linear(self.d_z, D_S) | |
| # EMA target encoder as buffers (saved in the state_dict, never in parameters(): not optimised, not counted as trainable) | |
| self.register_buffer("ema_w1", self.enc[0].weight.detach().clone()) | |
| self.register_buffer("ema_b1", self.enc[0].bias.detach().clone()) | |
| self.register_buffer("ema_w2", self.enc[2].weight.detach().clone()) | |
| self.register_buffer("ema_b2", self.enc[2].bias.detach().clone()) | |
| # ---------------------------------------------------------------- forward pieces | |
| def predict(self, z_v, h_u, z_ecat): | |
| x = torch.cat([z_v, h_u, z_ecat], 1) | |
| return self.pred(x).view(-1, self.K, self.d_z) | |
| def enc_input(s, v): | |
| return torch.cat([s.float() * v.float(), v.float()], -1) | |
| def encode_ema(self, x): | |
| with torch.no_grad(): | |
| h = F.gelu(F.linear(x, self.ema_w1, self.ema_b1)) | |
| return F.linear(h, self.ema_w2, self.ema_b2) | |
| def ema_update(self): | |
| m = self.m | |
| self.ema_w1.mul_(m).add_(self.enc[0].weight.detach(), alpha=1 - m) | |
| self.ema_b1.mul_(m).add_(self.enc[0].bias.detach(), alpha=1 - m) | |
| self.ema_w2.mul_(m).add_(self.enc[2].weight.detach(), alpha=1 - m) | |
| self.ema_b2.mul_(m).add_(self.enc[2].bias.detach(), alpha=1 - m) | |
| def variance_term(z, eps=1e-4): | |
| """VICReg variance regulariser on z [N, d]: mean_j relu(1 - sqrt(Var_j + eps)); 0 with fewer than 2 rows.""" | |
| if z.shape[0] < 2: | |
| return z.new_zeros(()) | |
| std = torch.sqrt(z.var(0, unbiased=False) + eps) | |
| return torch.relu(1.0 - std).mean() | |
| def loss(self, pred, tg): | |
| """pred [B,K,d_z] (any dtype; cast to float), tg = build_targets(...) -> (L_aux, parts lat_pred / lat_rec / lat_var / lat_n).""" | |
| s, v, ok = tg["s"].float(), tg["v"], tg["row_ok"] | |
| x = self.enc_input(s, v) | |
| z_on = self.enc(x) # [B,K,d_z] online (trainable by rec + var) | |
| z_tg = self.encode_ema(x) # [B,K,d_z] EMA target, no grad | |
| okf = ok.float() | |
| l_pred = (((pred.float() - z_tg) ** 2).mean(-1) * okf).sum() / okf.sum().clamp_min(1.0) | |
| rec = self.dec(z_on) | |
| vf = v.float() | |
| rec_c = ((rec[..., :N_CONT] - s[..., :N_CONT]) ** 2) * vf[..., :N_CONT] | |
| rec_b = F.binary_cross_entropy_with_logits(rec[..., N_CONT:], s[..., N_CONT:], reduction="none") * vf[..., N_CONT:] | |
| l_rec = (rec_c.sum() + rec_b.sum()) / vf.sum().clamp_min(1.0) | |
| l_var = self.variance_term(z_on[ok]) | |
| total = l_pred + self.rec_coef * l_rec + self.var_coef * l_var | |
| return total, {"lat_pred": l_pred, "lat_rec": l_rec, "lat_var": l_var, "lat_n": okf.sum().detach()} | |
| def n_params(m): | |
| return sum(p.numel() for p in m.parameters() if p.requires_grad) | |
Xet Storage Details
- Size:
- 10.7 kB
- Xet hash:
- b35f7c29472fb3322ad568f54f54ada10265c6ac169e377b0b0498214204d033
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.