Buckets:
| """cwm3/models/cwm3.py — the joint CWM3 model: three pretrained base branches + directed zero-init interactions + decoders (plan §3). | |
| V (VehicleBranch, V0/can4s | can8s) trunk z_v = TCN(X) [B,256] ⊕ mode(16) -> V emissions [B,25,35] | |
| U (DriverBranch, U0/kin_can_noK) trunk h_u = GRU(Kx)[-1] [B,128] ⊕ mode(16) -> U emissions [B,25,26] | |
| E (EnvBranch, E0/pooled | tok8s) tokens -> causal Transformer -> h_tok [B,17|65,128]; z_e = LN(h_tok[:, -1]) [B,128]; z_vf [B,128] | |
| ⊕ mode(16) -> E emissions [B,25,12], F [B,5,256] | |
| Interaction I (models/interaction.py; each residual is exactly 0 at init): | |
| road2drv : h_u += Road2Driver(h_u, keys) keys = RegionTokenEmbed(region token bank) [B,64,128] (tok_source region, default) | h_tok (enc) | |
| env2veh : z_v += Env2Vehicle([z_e ; z_vf]) | |
| drv2veh : z_v += gate(mode) * Driver2Vehicle(h_u) (uses the post-road2drv driver trunk) | |
| naive : CONTROL — [z_v ; h_u ; z_e ; z_vf] -> shared MLP (non-zero init) -> residuals on all three trunks. | |
| K residual head (k_head=True): a fresh Linear-GELU-Linear head on the DETACHED driver trunk [h_u ; mode] emits ΔK mu / log sig | |
| (U dims 0:22; K̂ = K_0 + ΔK̂). Detached -> its Gaussian loss cannot reach the GRU, so the hands_on / head_turned chains are | |
| trained exactly as in the chains-only base (kin_can_noK) while K̂ exists for the warning head. The chain dims 22:26 come from | |
| the pretrained U head. | |
| Detached auxiliary heads for the warning head (same pattern as the K head; no gradient into any trunk, so the acceptance comparisons are | |
| untouched): p_head = Linear-GELU-Linear on zE.detach() -> road-pressure chains P (ttc, ttc_drop, cut_in, lead_braking, lane_excursion, | |
| badas; results/targets_pressure) [B,25,12]; r_head on zV.detach() -> the HARD response chain (R >= 2; results/targets_response) [B,25,2]. | |
| Decoders (models/decoders.py): direct (pretrained flat heads; identity at init) | direct_fresh | chunked (parameter-matched). | |
| Gaze U bases (2026-09-12 appendix variant, recipes `u_gaze_chain` / `u_face` forwarded by _model_kwargs): the U base's gaze_head emits out["G"] [B,25,2] | |
| (gaze_off chain) on the same zU; its NLL joins the U loss with lambda G (default 1.0). Absent from every older config -> nothing changes. | |
| The forward takes the batch dicts of train/batch.py (b["V"] prep_V, b["U"] prep_U, b["E"] prep_E of the E base's token config, and b["R"] = | |
| prep_E region tokens when tok_source == "region") — nothing else; no target enters the model. No openpilot DMS field is read anywhere. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| from . import heads as hd | |
| from .branches import build_branch, flat_head, params_by_module | |
| from .decoders import EXPECT, build_decoder | |
| from .interaction import Driver2Vehicle, Env2Vehicle, NaiveFusion, Road2Driver, RegionTokenEmbed, n_params | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| ALL_INTERACTIONS = ("road2drv", "env2veh", "drv2veh") | |
| P_CHAINS = ("ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas") | |
| R_CHAINS = ("hard",) | |
| BASE_DEFAULT = {"V": ("V0", "can4s"), "U": ("U0", "kin_can_noK"), "E": ("E0", "pooled")} # every base is a CONFIG FIELD (stages_joint.yaml) | |
| D_MODE = 16 | |
| def n_face_channels(): | |
| """len(data/cwm3_dataset.FACE_CHANNELS) — the N_FACE of train_cwm3.model_kwargs, imported lazily so the model module does not import the data layer.""" | |
| d = str(CWM3 / "data") | |
| if d not in sys.path: | |
| sys.path.insert(0, d) | |
| import cwm3_dataset as ds | |
| return int(ds.N_FACE) | |
| def _model_kwargs(cfg): | |
| """train_cwm3.model_kwargs, duplicated here so the model module does not import the running trainer. | |
| 2026-09-11 (final assembly, roadmap §三 1): forwards the survivor-recipe switches — U `in_ln` (recipe key `u_in_ln`), E `d_side` (side vector | |
| `regionpca256` = 512) and E `p_head` (variants `*_P`) — each only when set, so every older config yields exactly the kwargs it did before.""" | |
| b, m = cfg["branch"], dict(cfg["model"]) | |
| rec = cfg.get("recipe", {}) or {} | |
| if cfg.get("arch"): | |
| # 2026-09-17 (owner experiment C): architecture-swap U bases (variant key `arch`, e.g. `driverwm` -> models/u_driverwm.py) are rebuilt with their own trunk | |
| # when a joint run loads them (train_cwm3.model_kwargs forwards the same key); absent from every older config -> nothing changes. | |
| m["arch"] = cfg["arch"] | |
| if b == "U": | |
| m["cabin"] = cfg.get("cabin", "none") | |
| if bool(rec.get("u_in_ln", False)): | |
| m["in_ln"] = True | |
| # 2026-09-12 (gaze-chain APPENDIX variant, coordinator assignment): the gaze U bases `*_gz` / `*_gzfc` — recipe `u_gaze_chain` -> DriverBranch | |
| # gaze_head (a third 2-state driver chain, output "G"), `u_face` -> n_face = the causal face-landmark channels appended to the GRU input by | |
| # train/batch.prep_U (the rule of train_cwm3.model_kwargs); both only when set, so every older config (incl. the frozen final bases) is unchanged. | |
| if bool(rec.get("u_gaze_chain", False)): | |
| m["gaze_chain"] = True | |
| if bool(rec.get("u_face", False)): | |
| m["n_face"] = n_face_channels() | |
| if cfg.get("n_kin"): | |
| m["n_kin"] = int(cfg["n_kin"]) # kin v2p (2026-09-18): 14 input kinematics columns; absent from every older config -> the default 11 | |
| if b == "E": | |
| m["tokens"], m["n_slots"] = cfg.get("tokens", "region"), int(cfg.get("n_slots", 16)) | |
| if cfg.get("side", "vf256") == "regionpca256": | |
| m["d_side"] = 512 | |
| if bool(cfg.get("p_head", False)): | |
| m["p_head"] = True | |
| return m | |
| BASE_CFG_KEYS = ("can_hist", "cabin", "tokens", "n_slots", "side", "p_head", "kclip", "kin_source", "kin_version", "stats_version", "pack_version", "mask_routes", | |
| "arch", "n_kin") # arch (2026-09-17, experiment C): recorded for architecture-swap bases; n_kin (2026-09-18, kin v2p): 14; absent from every older config -> stays absent | |
| def base_cfg_summary(c): | |
| """the base run's config -> the `base_cfgs` record of the joint run (base_ckpts.json / run.json) AND the per-branch batch.py cfg source: | |
| besides model / variant / lambda it carries every survivor-recipe switch (kclip -> batch.k_clip returns the base's clamp instead of the legacy 5, | |
| u_in_ln, kin_source / kin_version, stats_version, p_head, select_LF, side). Keys absent from an old config stay absent (old runs reproduce).""" | |
| out = {"model": c["model"], "variant": c["variant"], "lambda": c.get("lambda", {}), **{k: c[k] for k in BASE_CFG_KEYS if k in c}} | |
| rec = c.get("recipe", {}) or {} | |
| for k in ("u_in_ln", "select_LF", "u_gaze_chain", "u_face"): # u_gaze_chain / u_face (2026-09-12): read by _model_kwargs and the joint batch prep (prep_cfgs) | |
| if k in rec: | |
| out[k] = rec[k] | |
| return out | |
| def load_base_branch(run_dir, device="cpu"): | |
| """results/runs/<STAGE>_<variant>_s<k>/BEST -> (branch module with the BEST weights, cfg, stats dict, ckpt name).""" | |
| rd = Path(run_dir) | |
| best = (rd / "BEST").read_text().strip() | |
| z = torch.load(rd / best, map_location=device, weights_only=False) | |
| cfg = z["cfg"] | |
| m = build_branch(cfg["branch"], int(z["seed"]), **_model_kwargs(cfg)) | |
| m.load_state_dict(z["model"]) | |
| stats = {k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in z["stats"].items()} | |
| return m, cfg, stats, best | |
| def merge_stats(parts): | |
| """the three base runs' train-only statistics: overlapping keys must agree (same train split, same rule).""" | |
| import numpy as np | |
| out = {} | |
| for st in parts: | |
| for k, v in st.items(): | |
| v = np.asarray(v, np.float32) | |
| if k in out: | |
| assert np.allclose(out[k], v, rtol=1e-5, atol=1e-6), f"base statistics disagree on {k}" | |
| else: | |
| out[k] = v | |
| return out | |
| def stats_per_branch(parts, branches=("V", "U", "E")): | |
| """stats_mode `per_branch` (roadmap §三 3, explicit CLI choice — the merged-mode 1e-5 assert above is never loosened): every base keeps its own | |
| statistics under the keys '<branch>/<key>'; train_cwm3_joint.prep_joint hands each batch.prep_* the sub-dict of its branch (branch_stats).""" | |
| import numpy as np | |
| out = {} | |
| for br, st in zip(branches, parts): | |
| for k, v in st.items(): | |
| out[f"{br}/{k}"] = np.asarray(v, np.float32) | |
| return out | |
| def branch_stats(st, br): | |
| """the statistics dict a branch's batch.prep_* sees: merged mode (no '<br>/' keys) -> st itself; per-branch mode -> its '<br>/<k>' keys renamed | |
| to <k> plus every unprefixed key (the auxiliary base rates p0_base / r0_base, added from the TRAIN files).""" | |
| pre = f"{br}/" | |
| if not any(str(k).startswith(pre) for k in st): | |
| return st | |
| out = {str(k): v for k, v in st.items() if "/" not in str(k)} | |
| out.update({str(k)[len(pre):]: v for k, v in st.items() if str(k).startswith(pre)}) | |
| return out | |
| class CWM3Joint(nn.Module): | |
| def __init__(self, V, U, E, interactions=ALL_INTERACTIONS, naive=False, decoder="direct", k_head=True, i_hidden=640, dropout=0.1, | |
| k_hidden=128, seed=0, p_head=True, r_head=True, tok_source="region", p_head_from_base=True, latent_aux=None, kpf_dim=0, mask_V=False): | |
| """mask_V (2026-09-18, owner question 'front video + driver only, CAN missing'; pre-registration docs/decisions.md 2026-09-18T03:34:28Z, variant `ue_noV`): | |
| the vehicle branch is treated as a MISSING modality — trunk_V returns EXACT zeros in place of z_v (the way vf_ok / cab_ok zero a missing side input), so no | |
| V-derived quantity reaches any output; the V decoder / R head then see [0 ; mode] only (their losses are masked out by the trainer's all-False V-valid mask). | |
| Default False = every older model bit for bit (tests/test_mask_v.py). | |
| latent_aux (2026-09-17, R2 WP4 batch 2 C4; models/latent_aux.py): None (default, every older config) | dict of LatentAux kwargs -> the low-dimensional | |
| multi-horizon latent auxiliary predictor on the post-interaction trunks [z_v ; h_u ; z_e ; z_vf]; out["LAT"] [B,K,d_z]; its loss is added by the trainer. | |
| tok_source: keys / values of the road -> driver cross-attention — "region": the raw region token bank [B,16,4,64] through a light | |
| embedding (RegionTokenEmbed; the E base may be pooled, its chains / F stay what B2 compares against); "enc": the E encoder's own | |
| token states h_tok (region or pooled, contextualised). | |
| p_head_from_base (2026-09-11, roadmap §三 2): when the E base carries a trunk-supervised P head of the same shape (E0 `*_P` variants), | |
| the joint's detached P head starts as a copy of it (P emissions identical at init, like the other heads); otherwise it is a fresh head | |
| (recorded in `p_head_source`, e.g. the E3 `no P` ablation whose E base has no P head).""" | |
| super().__init__() | |
| torch.manual_seed(seed) # new modules (I, K head, fresh decoders) seeded; base weights untouched | |
| self.V, self.U, self.E = V, U, E | |
| self.mask_V = bool(mask_V) | |
| self.interactions = tuple(x for x in interactions if x) if not naive else () | |
| self.naive = naive | |
| self.decoder_kind = decoder | |
| self.tok_source = tok_source | |
| # d_u: the GRU width for every reference base; `d_trunk` for architecture-swap U bases that expose encode_trunk (models/u_driverwm.py, 2026-09-17 experiment C) | |
| d_v, d_u, d_e = V.tcn.proj.out_channels, (U.gru.hidden_size if hasattr(U, "gru") else int(U.d_trunk)), E.norm.normalized_shape[0] | |
| d_vf = E.vf[1].out_features | |
| self.d_cab = U.cab[1].out_features if U.cab is not None else 0 | |
| self.road2drv = Road2Driver(d_u, d_e, d_e, 4, i_hidden, dropout) if "road2drv" in self.interactions else None | |
| self.tok_embed = RegionTokenEmbed(64, d_e, 16, 4) if (self.road2drv is not None and tok_source == "region") else None | |
| self.env2veh = Env2Vehicle(d_e + d_vf, d_v, i_hidden, dropout) if "env2veh" in self.interactions else None | |
| self.drv2veh = Driver2Vehicle(d_u, d_v, i_hidden, dropout) if "drv2veh" in self.interactions else None | |
| self.fusion = NaiveFusion((d_v, d_u, d_e + d_vf), i_hidden, dropout) if naive else None | |
| # emission decoders on the (possibly interaction-augmented) trunks | |
| zV, zU, zE = d_v + D_MODE, d_u + D_MODE + self.d_cab, d_e + d_vf + D_MODE | |
| hV, hU, hE = V.head[0].out_features, U.head[0].out_features, E.head[0].out_features | |
| dV, dU, dE = (V.head[2].p if isinstance(V.head[2], nn.Dropout) else 0.0), (U.head[2].p if isinstance(U.head[2], nn.Dropout) else 0.0), \ | |
| (E.head[2].p if isinstance(E.head[2], nn.Dropout) else 0.0) | |
| self.dec_V, iV = build_decoder(decoder, "V", zV, hV, dV, pretrained_net=V.head) | |
| self.dec_U, iU = build_decoder(decoder, "U", zU, hU, dU, pretrained_net=U.head) | |
| self.dec_E, iE = build_decoder(decoder, "E", zE, hE, dE, pretrained_net=E.head) | |
| self.decoder_info = {"V": iV, "U": iU, "E": iE} | |
| self.k_head = flat_head(d_u + D_MODE + self.d_cab, k_hidden, hd.H * 22, 0.0) if k_head else None | |
| self.p_head = flat_head(zE, hE, hd.H * 2 * len(P_CHAINS), dE) if p_head else None # detached road-pressure chains | |
| self.r_head = flat_head(zV, 128, hd.H * 2 * len(R_CHAINS), dV) if r_head else None # detached hard-response chain | |
| self.p_head_source = "none" if self.p_head is None else "fresh (random init; the E base has no P head or shapes differ)" | |
| self.latent_aux = None | |
| if latent_aux: | |
| from .latent_aux import LatentAux, KW as LAT_KW | |
| self.latent_aux = LatentAux(d_v + d_u + d_e + d_vf, **{k: v for k, v in dict(latent_aux).items() if k in LAT_KW}) | |
| e_p = getattr(E, "p_head", None) | |
| if self.p_head is not None and e_p is not None and p_head_from_base: | |
| src, dst = e_p.state_dict(), self.p_head.state_dict() | |
| if set(src) == set(dst) and all(src[k].shape == dst[k].shape for k in dst): | |
| self.p_head.load_state_dict(src) | |
| self.p_head_source = "copied from the E base's trunk-supervised P head (identity at init)" | |
| # LANE F1 (2026-09-17, pre-registration docs/decisions.md 2026-09-17T19:33:12Z): the DETACHED keypoint-forecaster channel (models/kpf_inject.py) as a | |
| # zero-initialised residual on h_u; created LAST so every module above keeps its paired seed init; kpf_dim 0 (default) -> bitwise the old model. | |
| self.kpf = None | |
| if int(kpf_dim) > 0: | |
| from .kpf_inject import KpfInject | |
| self.kpf = KpfInject(int(kpf_dim), d_u) | |
| # ---------------------------------------------------------------- trunks (the base forwards, split before their heads) | |
| def trunk_V(self, b): | |
| if self.mask_V: | |
| # missing vehicle modality (variant `ue_noV`, 2026-09-18): exact zeros of the trunk width, no TCN forward — nothing of X can enter any output | |
| x = b["X"] | |
| return x.new_zeros((x.shape[0], self.V.tcn.proj.out_channels), dtype=torch.float32), self.V.mode(b["mode"].long()) | |
| return self.V.tcn(b["X"].float()), self.V.mode(b["mode"].long()) | |
| def trunk_U(self, b): | |
| kx = b["Kx"].float() | |
| if getattr(self.U, "in_ln", None) is not None: # survivor recipe `u_in_ln` (2026-09-11): the base's input LayerNorm over the 12 kinematics channels | |
| kx = torch.cat([self.U.in_ln(kx[..., :self.U.N_KIN_IN]), kx[..., self.U.N_KIN_IN:]], -1) | |
| if getattr(self.U, "encode_trunk", None) is not None: | |
| h_last = self.U.encode_trunk(kx) # architecture-swap U trunk (models/u_driverwm.py, 2026-09-17 experiment C): same LN-applied Kx in, [B,d_u] out | |
| else: | |
| _, h = self.U.gru(kx) | |
| h_last = h[-1] | |
| extra = [self.U.mode(b["mode"].long())] | |
| if self.U.cab is not None: | |
| extra.append(self.U.cab(b["cab"].float()) * b["cab_ok"].float()[:, None]) | |
| return h_last, extra | |
| def trunk_E(self, b): | |
| E = self.E | |
| tok = b["tok"].float() | |
| B, L, R, _ = tok.shape | |
| assert L == E.n_slots and R == E.n_reg | |
| x = E.tok_in(tok) | |
| miss = (~b["tok_ok"].bool())[:, :, None, None].float() | |
| x = x * (1 - miss) + E.missing * miss | |
| x = x + E.slot_emb.weight[None, :, None, :] + E.reg_emb.weight[None, None, :, :] | |
| x = x + E.dt_in(b["tok_dt"].float()[:, :, None, None]) | |
| x = x.reshape(B, L * R, -1) | |
| x = torch.cat([x, E.readout[None, None].expand(B, 1, -1)], 1) | |
| h_tok = E.enc(x, mask=E.attn_mask) | |
| z_e = E.norm(h_tok[:, -1]) | |
| z_vf = E.vf(b["vf"].float()) * b["vf_ok"].float()[:, None] | |
| return h_tok, z_e, z_vf, E.mode(b["mode"].long()) | |
| # ---------------------------------------------------------------- joint forward | |
| def forward(self, b, return_trunks=False): | |
| bV, bU, bE = b["V"], b["U"], b["E"] | |
| z_v, m_v = self.trunk_V(bV) | |
| h_u, extra_u = self.trunk_U(bU) | |
| if self.kpf is not None: | |
| h_u = h_u + self.kpf(bU["kpf"]) # LANE F1 injector (zero at init); bU["kpf"] = the forecaster channel of train/batch.prep_U | |
| h_tok, z_e, z_vf, m_e = self.trunk_E(bE) | |
| z_ecat = torch.cat([z_e, z_vf], 1) | |
| if self.road2drv is not None: | |
| keys = self.tok_embed(b["R"]) if self.tok_embed is not None else h_tok | |
| h_u = h_u + self.road2drv(h_u, keys) | |
| if self.env2veh is not None: | |
| z_v = z_v + self.env2veh(z_ecat) | |
| if self.drv2veh is not None: | |
| z_v = z_v + self.drv2veh(h_u, bV["mode"]) | |
| if self.fusion is not None: | |
| r_v, r_u, r_e = self.fusion(z_v, h_u, z_ecat) | |
| z_v, h_u, z_ecat = z_v + r_v, h_u + r_u, z_ecat + r_e | |
| zV = torch.cat([z_v, m_v], 1) | |
| zU = torch.cat([h_u] + extra_u, 1) | |
| zE = torch.cat([z_ecat, m_e], 1) | |
| V = self.dec_V(zV, EXPECT["V"][0]) | |
| U = self.dec_U(zU, EXPECT["U"][0]) | |
| E = self.dec_E(zE, EXPECT["E"][0]) | |
| Fh = self.E.f_head(zE).view(-1, 5, hd.D_F) | |
| if self.k_head is not None: | |
| K = self.k_head(zU.detach()).view(-1, hd.H, 22) | |
| U = torch.cat([K, U[..., 22:]], -1) | |
| out = {"V": V, "U": U, "E": E, "F": Fh} | |
| if getattr(self.U, "gaze_head", None) is not None: | |
| out["G"] = self.U.gaze_head(zU).view(-1, hd.H, 2) # gaze_off chain of the gaze U bases (2026-09-12): the base's head on the same z = [h_u ; mode (; cab)] | |
| if self.p_head is not None: | |
| out["P"] = self.p_head(zE.detach()).view(-1, hd.H, 2 * len(P_CHAINS)) | |
| if self.r_head is not None: | |
| out["R"] = self.r_head(zV.detach()).view(-1, hd.H, 2 * len(R_CHAINS)) | |
| if self.latent_aux is not None: | |
| out["LAT"] = self.latent_aux.predict(z_v, h_u, z_ecat) # C4 (2026-09-17): NOT detached — the auxiliary loss reaches the trunks by design | |
| if return_trunks: | |
| out["trunks"] = {"z_v": z_v, "h_u": h_u, "z_e": z_e, "z_vf": z_vf, "h_tok": h_tok} | |
| return out | |
| # ---------------------------------------------------------------- parameter groups / reports | |
| def base_parameters(self): | |
| for m in (self.V, self.U, self.E): | |
| yield from m.parameters() | |
| def new_parameters(self): | |
| base = {id(p) for p in self.base_parameters()} | |
| return [p for p in self.parameters() if id(p) not in base] | |
| def param_report(self): | |
| rep = {"V": params_by_module(self.V), "U": params_by_module(self.U), "E": params_by_module(self.E)} | |
| for name in ("road2drv", "env2veh", "drv2veh", "fusion", "k_head", "p_head", "r_head", "tok_embed", "kpf"): | |
| m = getattr(self, name) | |
| rep[name] = n_params(m) if m is not None else 0 | |
| rep["latent_aux"] = n_params(self.latent_aux) if self.latent_aux is not None else 0 # C4 trainable (the EMA target encoder is buffers, not parameters) | |
| rep["latent_aux_ema_buffers"] = int(sum(b.numel() for b in self.latent_aux.buffers())) if self.latent_aux is not None else 0 | |
| rep["decoders"] = self.decoder_info | |
| rep["interaction_total"] = sum(rep[n] for n in ("road2drv", "env2veh", "drv2veh", "fusion", "tok_embed")) | |
| rep["new_total"] = sum(p.numel() for p in self.new_parameters()) | |
| rep["total"] = sum(p.numel() for p in self.parameters()) | |
| return rep | |
| def build_joint(seed, base=BASE_DEFAULT, runs_dir=CWM3 / "results" / "runs", device="cpu", dirs=None, stats_mode="merged", **kw): | |
| """Assemble the joint model from the seed-k BEST checkpoints of the three base stages. -> (model, info). | |
| dirs: optional {branch: run dir} overriding the default `<stage>_<variant>_s<seed>` names (needed for suffixed bases such as `_v2`; | |
| 2026-09-11 bug: without it the v2 joint runs loaded the v1 E weights while the parent reported the v2 names). | |
| stats_mode: merged (default; overlapping statistics asserted equal to 1e-5) | per_branch (each branch its own stats.npz, keys '<br>/<k>').""" | |
| mods, cfgs, stats, names = {}, {}, [], {} | |
| for br, (stage, variant) in base.items(): | |
| rd = Path(dirs[br]) if dirs and br in dirs else Path(runs_dir) / f"{stage}_{variant}_s{seed}" | |
| assert (rd / "BEST").exists() and (rd / "TRAIN_DONE").exists(), f"base checkpoint missing: {rd}" | |
| m, cfg, st, best = load_base_branch(rd, device) | |
| mods[br], cfgs[br], names[br] = m, cfg, f"{rd.name}:{best}" | |
| stats.append(st) | |
| model = CWM3Joint(mods["V"], mods["U"], mods["E"], seed=seed, **kw) | |
| info = {"base_ckpts": names, "base_cfgs": {k: base_cfg_summary(v) for k, v in cfgs.items()}, | |
| "stats": merge_stats(stats) if stats_mode == "merged" else stats_per_branch(stats, tuple(base)), "stats_mode": stats_mode, | |
| "params": model.param_report(), "p_head_source": model.p_head_source} | |
| return model, info | |
| def aux_chain_loss(logits, tg, names, prefix): | |
| """detached auxiliary chains: logits [B,T,2C]; tg: z_prev / z / m [B,T,C]. Parts are named <prefix>_<chain> (NOT chain_*: excluded from the | |
| selection objective — the heads are detached, their NLL says nothing about the trunks).""" | |
| c = logits.float().reshape(*logits.shape[:-1], len(names), 2) | |
| parts, total = {}, 0.0 | |
| for j, name in enumerate(names): | |
| l = hd.masked_mean(hd.chain_bce(c[..., j, 0], c[..., j, 1], tg["z_prev"][..., j], tg["z"][..., j], tg["m"][..., j]), tg["m"][..., j]) | |
| parts[f"{prefix}_{name}"] = l | |
| total = total + l | |
| return total, parts | |
| def joint_loss(out, tg, lam, return_branches=False): | |
| """Sum of the three branch losses (heads.py forms) + the detached auxiliary chains; tg = {"V": tgV, "U": tgU, "E": tgE[, "P": tgP, "R": tgR]}. | |
| -> total, parts (chain_* + L_F form the selection objective; pchain_* / rchain_* do not). | |
| return_branches=True (2026-09-16, WP3 MTL block; train/mtl.py) -> total, parts, branches with branches = {"V": lv, "U": lu, "E": le, "aux": the | |
| detached P / R terms} — the SAME tensors that `total` sums (total is computed exactly as before: callers with the two-value contract are unchanged).""" | |
| lv, pv = hd.vehicle_loss(out["V"], tg["V"], lam_hz=lam.get("hz", 0.5)) | |
| lu, pu = hd.driver_loss(out["U"], tg["U"], lam_K=lam.get("K", 0.5)) | |
| if "G" in out and "zG" in tg["U"]: | |
| # gaze_off chain (2026-09-12, gaze U bases): lambda G (default 1.0 = the U base's) x its NLL; the part chain_gaze_off enters the selection objective like | |
| # every other chain (heads.selection_objective) — as in the U base run. Without "G" or its targets (every older config): exactly driver_loss. | |
| lg, pg = hd.gaze_loss(out["G"], tg["U"]) | |
| pu.update(pg) | |
| lu = lu + float(lam.get("G", 1.0)) * lg | |
| le, pe = hd.env_loss(out["E"], out["F"], tg["E"], lam_F=lam.get("F", 0.1)) | |
| parts = {} | |
| parts.update(pv) | |
| parts.update(pu) | |
| parts.update(pe) | |
| total = lv + lu + le | |
| laux = 0.0 | |
| if "P" in out and "P" in tg: | |
| lp, pp = aux_chain_loss(out["P"], tg["P"], P_CHAINS, "pchain") | |
| parts.update(pp) | |
| total = total + lam.get("P", 1.0) * lp | |
| laux = laux + lam.get("P", 1.0) * lp | |
| if "R" in out and "R" in tg: | |
| lr, prr = aux_chain_loss(out["R"], tg["R"], R_CHAINS, "rchain") | |
| parts.update(prr) | |
| total = total + lam.get("R", 1.0) * lr | |
| laux = laux + lam.get("R", 1.0) * lr | |
| if return_branches: | |
| if not torch.is_tensor(laux): | |
| laux = torch.zeros((), dtype=total.dtype, device=total.device) | |
| return total, parts, {"V": lv, "U": lu, "E": le, "aux": laux} | |
| return total, parts | |
| if __name__ == "__main__": | |
| m, info = build_joint(0) | |
| print(json.dumps(info["params"], indent=1)) | |
| print("base:", info["base_ckpts"]) | |
Xet Storage Details
- Size:
- 25.1 kB
- Xet hash:
- 4e39e8a3ae88d5b0d2095f5c8201946678235f89c39feb0ddf18628c878505b0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.