Buckets:
| """cwm3/models/future_state_head.py — F2′ (R2 WP4 batch 1, 2026-09-16; plan <local-notes>/plans/elegant-dazzling-stallman.md; pre-registration docs/decisions.md | |
| 2026-09-16T06:46:35Z): a DETACHED, horizon-conditioned future-state head on the frozen joint model's driver state. | |
| Inputs (per anchor t, nothing after t): | |
| z_U [B,128] the frozen FINAL joint model's driver trunk state h_u (train/train_future_state.py recomputes it from the U stream; the model has no | |
| interaction INTO the driver branch, so h_u = GRU(Kx) of the fine-tuned U trunk — verified equal to the saved prediction pack) | |
| dwell [B,D] label-free dwell / switch / observation-confidence features of the 8 s history (train/train_future_state.FEATURES; standardised, NaN -> 0 | |
| + missing flags): time in the current state, time since the last switch, switch counts, validity shares, current gaze state / dwell | |
| mode [B] regime (0 manual / 1 engaged) | |
| z0 [B,S] the CURRENT state of each chain variable (hands_on, head_turned, gaze_off) = pi_0 of the recursion | |
| Outputs per state s (S = 3; the states CO-OCCUR -> multi-task, NOT competing risks): | |
| chains [B,H,S,2] (h_on, h_off) logits per horizon k = 1..H (0.2 s steps, H = 25 = 5 s) — bidirectional transition hazards | |
| pi [B,H,S] occupancy pi_k = pi_{k-1} (1 - h_off_k) + (1 - pi_{k-1}) h_on_k, pi_0 = z0 (models/heads.forward_marginal, the FINAL model's recursion) | |
| first_leave [B,S,H+1] logits of the discrete first-leave-time law: bin k-1 = the first step k at which the state differs from z0, bin H = "no leave within H" | |
| rem_dwell [B,S,H+1] the same law, trained on rows whose CURRENT state is the off-task state only (hands_on = 0 / head_turned = 1 / gaze_off = 1): | |
| remaining dwell in the current off state (= recovery time) | |
| Derived (future_state_outputs): E[D_off] = dt * sum_k pi_off_k (expected cumulative off-task seconds over the horizon), E[min(T, (H+1) dt)] of the two laws, | |
| P(leave within 1 / 2 / 3 / 5 s). | |
| Losses (future_state_loss): transition NLL on the observed switches (models/heads.chain_bce — the existing chain form, masks of train/batch.prep_U) + | |
| masked BCE(pi_k, s_{t+k}) (calibration first; focal / ranking losses are later controls) + discrete-time survival NLL of the first-leave and | |
| remaining-dwell laws with right censoring (censored at the last observed step c: -log sum_{j>c} p_j). | |
| Strict future: every target index k-1 refers to the state at t + k dt, k >= 1 (the kin / gaze target contract S_c / S_g); tests/test_future_state_head.py | |
| poisons the future and checks the inputs do not move. ~0.035 M parameters (n_params()). No openpilot DMS field anywhere (states = our own pose / | |
| landmark chains). | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from . import heads as hd | |
| H = hd.H # 25 steps | |
| DT = 0.2 # s per step | |
| STATES = ("hands_on", "head_turned", "gaze_off") | |
| OFF_VALUE = {"hands_on": 0.0, "head_turned": 1.0, "gaze_off": 1.0} # the off-task value of each chain variable | |
| N_BINS = H + 1 # first-leave bins: step 1..H, then "beyond H" | |
| LEAVE_TAUS = {"1s": 5, "2s": 10, "3s": 15, "5s": 25} | |
| class FutureStateHead(nn.Module): | |
| def __init__(self, d_z=128, d_dwell=19, n_states=len(STATES), d=64, d_h=16, dropout=0.0): | |
| super().__init__() | |
| self.n_states, self.d = int(n_states), int(d) | |
| self.z_ln = nn.LayerNorm(d_z) | |
| self.z_in = nn.Linear(d_z, d) | |
| self.dw_in = nn.Linear(d_dwell, d) | |
| self.mode = nn.Embedding(2, d) | |
| self.hor = nn.Embedding(H, d_h) | |
| self.trans = nn.Sequential(nn.Linear(d + d_h, d), nn.GELU(), nn.Dropout(dropout), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 2 * self.n_states)) | |
| self.row = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Dropout(dropout), nn.Linear(d, 2 * self.n_states * N_BINS)) | |
| def context(self, z, dwell, mode): | |
| x = self.z_in(self.z_ln(z.float())) + self.dw_in(dwell.float()) + self.mode(mode.long()) | |
| return F.gelu(x) | |
| def forward(self, z, dwell, mode): | |
| B = z.shape[0] | |
| x = self.context(z, dwell, mode) # [B,d] | |
| xh = torch.cat([x[:, None, :].expand(-1, H, -1), self.hor.weight[None].expand(B, -1, -1)], -1) | |
| chains = self.trans(xh).view(B, H, self.n_states, 2) | |
| row = self.row(x).view(B, 2, self.n_states, N_BINS) | |
| return {"chains": chains, "first_leave": row[:, 0], "rem_dwell": row[:, 1]} | |
| def n_params(m): | |
| return int(sum(p.numel() for p in m.parameters())) | |
| # ----------------------------------------------------------------------------- occupancy / laws | |
| def occupancy(chains, z0): | |
| """chains [B,H,S,2] logits, z0 [B,S] current state in {0,1} -> pi [B,H,S] by the FINAL model's forward recursion (heads.forward_marginal per state).""" | |
| h_on, h_off = hd.hazards(chains) | |
| return torch.stack([hd.forward_marginal(h_on[:, :, s], h_off[:, :, s], z0[:, s]) for s in range(chains.shape[2])], -1) | |
| def leave_law(logits): | |
| """[.., N_BINS] logits -> probabilities p_k (k = 1..H) and p_beyond (index H).""" | |
| return torch.softmax(logits.float(), -1) | |
| def expected_leave_s(p): | |
| """E[min(T, (H+1) dt)] in seconds from a first-leave law p [.., N_BINS] (bin j -> step j+1; the 'beyond H' bin counted at step H+1).""" | |
| steps = torch.arange(1, N_BINS + 1, device=p.device, dtype=p.dtype) * DT | |
| return (p * steps).sum(-1) | |
| def leave_within(p, tau_steps): | |
| """P(first leave within tau_steps) = sum of the first tau_steps bins.""" | |
| return p[..., :tau_steps].sum(-1) | |
| def leave_targets(z0, S, valid): | |
| """First-leave target of one state: z0 [B] current state, S [B,H] future states (any value where ~valid), valid [B,H] observed steps. | |
| Observation prefix = the steps before the first unobserved one. -> (event [B] bool, k [B] long): event -> first differing step k in 1..H; | |
| no event -> k = number of observed prefix steps c (c == H: fully observed, class 'beyond H'; c < H: right-censored at c). Rows with c == 0 and no event | |
| are censored at 0 (loss = -log 1 = 0 -> excluded by the caller's mask).""" | |
| B, T = S.shape | |
| obs = valid.bool() | |
| prefix = torch.cumprod(obs.long(), 1) # 1 while every step so far is observed | |
| diff = (S.float() != z0.float()[:, None]) & (prefix > 0) | |
| has = diff.any(1) | |
| first = torch.where(has, diff.float().argmax(1) + 1, torch.zeros(B, dtype=torch.long, device=S.device)) | |
| c = prefix.sum(1) | |
| return has, torch.where(has, first, c) | |
| def leave_nll(logits, event, k, m): | |
| """Discrete-time survival NLL of a first-leave law with right censoring (per row, zero where m == 0). | |
| event -> -log p_k (bin k-1); no event, k == H -> -log p_beyond; no event, k < H -> -log sum_{j > k} p_j (bins k..H).""" | |
| lp = torch.log_softmax(logits.float(), -1) # [B,N_BINS] | |
| B = lp.shape[0] | |
| idx = (k - 1).clamp(0, H - 1) | |
| l_event = -lp.gather(1, idx[:, None])[:, 0] | |
| tail = torch.flip(torch.logcumsumexp(torch.flip(lp, [1]), 1), [1]) # tail[j] = log sum_{i >= j} p_i | |
| j = k.clamp(0, H) # censored at k: survive bins 0..k-1 -> mass in bins >= k | |
| l_cens = -tail.gather(1, j[:, None])[:, 0] | |
| l = torch.where(event.bool(), l_event, l_cens) | |
| return l * m.float() | |
| # ----------------------------------------------------------------------------- loss | |
| def future_state_loss(out, tg, lam_occ=1.0, lam_leave=1.0): | |
| """out = FutureStateHead.forward; tg per state s (stacked on the last axis, [B,H,S] / [B,S]): | |
| z_prev, z, m transition targets / mask (train/batch.prep_U conventions: z_prev = anchor state then z[:-1]; m = onset mask & M & finite) | |
| Sfut, Sm future state and its observation mask for the occupancy BCE (finite & M) | |
| z0, z0valid the current state and its validity (pi_0; rows with an invalid anchor are excluded from every term of that state) | |
| -> total, parts {trans_<s>, occ_<s>, leave_<s>, rem_<s>, occ_bce (selection objective)}.""" | |
| chains = out["chains"].float() | |
| S = chains.shape[2] | |
| z0v = tg["z0valid"].float() | |
| parts, total = {}, 0.0 | |
| pi = occupancy(chains, tg["z0"].float()) | |
| for s in range(S): | |
| name = STATES[s] if s < len(STATES) else f"s{s}" | |
| mt = tg["m"][..., s].float() * z0v[:, s:s + 1] | |
| lt = hd.masked_mean(hd.chain_bce(chains[:, :, s, 0], chains[:, :, s, 1], tg["z_prev"][..., s], tg["z"][..., s], mt), mt) | |
| mo = tg["Sm"][..., s].float() * z0v[:, s:s + 1] | |
| lo = hd.masked_mean(F.binary_cross_entropy(pi[:, :, s].clamp(1e-6, 1 - 1e-6), tg["Sfut"][..., s].float(), reduction="none"), mo) | |
| ev, k = leave_targets(tg["z0"][:, s], tg["Sfut"][..., s], tg["Sm"][..., s]) | |
| ml = z0v[:, s] * (ev | (k > 0)).float() | |
| ll = hd.masked_mean(leave_nll(out["first_leave"][:, s], ev, k, ml), ml) | |
| off = (tg["z0"][:, s].float() == OFF_VALUE.get(name, 1.0)).float() | |
| mr = ml * off | |
| lr = hd.masked_mean(leave_nll(out["rem_dwell"][:, s], ev, k, mr), mr) | |
| parts.update({f"trans_{name}": lt, f"occ_{name}": lo, f"leave_{name}": ll, f"rem_{name}": lr}) | |
| total = total + lt + lam_occ * lo + lam_leave * (ll + lr) | |
| parts["occ_bce"] = sum(parts[f"occ_{STATES[s] if s < len(STATES) else f's{s}'}"] for s in range(S)) / S | |
| return total, parts | |
| def future_state_outputs(out, z0): | |
| """-> dict of derived quantities: pi [B,H,S], pi_off [B,H,S], ED_off_s [B,S], h_on / h_off [B,H,S], first_leave_p / rem_dwell_p [B,S,N_BINS], | |
| first_leave_E_s / rem_dwell_E_s [B,S], leave_within_<tau> [B,S].""" | |
| chains = out["chains"].float() | |
| S = chains.shape[2] | |
| h_on, h_off = hd.hazards(chains) | |
| pi = occupancy(chains, z0.float()) | |
| offv = torch.tensor([OFF_VALUE.get(STATES[s] if s < len(STATES) else "", 1.0) for s in range(S)], device=pi.device) | |
| pi_off = torch.where(offv[None, None, :] > 0.5, pi, 1 - pi) | |
| fl, rd = leave_law(out["first_leave"]), leave_law(out["rem_dwell"]) | |
| res = {"pi": pi, "pi_off": pi_off, "ED_off_s": DT * pi_off.sum(1), "h_on": h_on, "h_off": h_off, "first_leave_p": fl, "rem_dwell_p": rd, | |
| "first_leave_E_s": expected_leave_s(fl), "rem_dwell_E_s": expected_leave_s(rd)} | |
| for nm, tau in LEAVE_TAUS.items(): | |
| res[f"leave_within_{nm}"] = leave_within(fl, tau) | |
| res[f"rem_within_{nm}"] = leave_within(rd, tau) | |
| return res | |
Xet Storage Details
- Size:
- 10.6 kB
- Xet hash:
- c305852d902d471fc4d0bc0b52d8bc6c379e36278b27bd43e727f6195fb117dd
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.