Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /models /overlap_head.py

AsianPlayer

16 days ago

download

raw

9.14 kB

"""cwm3/models/overlap_head.py — F8 (R2 WP4 batch 2, 2026-09-17; plan 'F8 驾驭员状态 × 道路需求的时间重叠预测'; pre-registration docs/decisions.md
2026-09-17T19:32:07Z): a DETACHED, horizon-conditioned JOINT head for P(driver unmet state U_{t+k}, road demand D_{t+k} | H_t) on the frozen ECAN joint model.
Variables per horizon k = 1..K (0.2 s steps, K = 15 = 3 s), all label-free target states of the existing packs (nothing after t enters the inputs):
  hands_off_{t+k}   = S_c[:, k-1, 0] == 0   (results/targets_kin_causal_routes_v2; NaN = unobserved)
  head_turned_{t+k} = S_c[:, k-1, 1] == 1
  D_{t+k}           = any road-pressure family (P ∧ Pvalid)[:, k-1] (results/targets_pressure, whitelist CAN margins); observed where any family is valid ∧ M
The head emits an 8-way categorical over (hands_off, head_turned, D) per horizon: state index s = 4·hands_off + 2·head_turned + D.  The four (U, D)
joint states and q_k = P(U = 1, D = 1) are marginalisations (overlap_from_joint): U = hands_off ∨ head_turned (manual, paper-facing) or U = head_turned
only (engaged paper-facing: the engaged hand cells are out of the paper and masked from every paper-facing endpoint).
Inputs: z [B, d_z] = LN([h_u ; z_e ; z_vf]) of the frozen joint model (detached cache), anchor [B, d_a] = label-free current states with validity
(hands_on0, head0, 6 P0 families + their validity bits), mode [B] regime, optional zk [B, K, d_k] per-horizon features (the `margonly` attribution arm:
the frozen model's own marginals, WITHOUT trunk features; use_z=False).
Baseline (pre-registered): the PRODUCT of the frozen model's own marginals — joint_from_marginals(p_hands_off, p_head, p_D) is the independent 8-way
distribution; overlap_from_joint of it equals p_U · p_D EXACTLY (identity unit test, tests/test_overlap_head.py).
Loss: masked joint NLL = −log Σ_{s consistent with the observed variables} p_s (the categorical form of a masked BCE: unobserved variables are
marginalised; a (row, k) with nothing observed contributes 0).  J_t = Δt Σ_k w_k q_k (w_k = 1 pre-fixed).  ≈ 0.06 M parameters (n_params()).
No openpilot DMS field anywhere (states = our own pose chains and CAN margins).
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F

K = 15                                        # horizons (0.2 s steps) = 3 s
DT = 0.2
VARS = ("hands_off", "head_turned", "D")      # bit order of the 8-way state index: s = 4*hands_off + 2*head_turned + D
N_STATES = 8
REPORT_STEPS = {"0.6s": 3, "1s": 5, "2s": 10, "3s": 15}   # 0.5 s is not on the grid: step 3 (0.6 s) = nearest step >= 0.5 s
D_ANCHOR = 2 + 2 + 6 + 6                      # hands_on0, head0, their validity, P0 x6, P0valid x6

# state bits [8, 3]: bit_v(s)
_BITS = torch.tensor([[(s >> 2) & 1, (s >> 1) & 1, s & 1] for s in range(N_STATES)], dtype=torch.float32)


def state_bits(device=None, dtype=torch.float32):
    b = _BITS.to(dtype)
    return b.to(device) if device is not None else b


class OverlapHead(nn.Module):
    def __init__(self, d_z=384, d_anchor=D_ANCHOR, d=96, d_h=16, n_h=K, d_k=0, use_z=True, dropout=0.0):
        super().__init__()
        self.use_z, self.n_h, self.d_k = bool(use_z), int(n_h), int(d_k)
        self.z_ln = nn.LayerNorm(d_z) if use_z else None
        self.z_in = nn.Linear(d_z, d) if use_z else None
        self.a_in = nn.Linear(d_anchor, d)
        self.mode = nn.Embedding(2, d)
        self.hor = nn.Embedding(n_h, d_h)
        self.net = nn.Sequential(nn.Linear(d + d_h + self.d_k, d), nn.GELU(), nn.Dropout(dropout), nn.Linear(d, d), nn.GELU(), nn.Linear(d, N_STATES))

    def forward(self, z, anchor, mode, zk=None):
        """-> logits [B, n_h, 8]."""
        x = self.a_in(anchor.float()) + self.mode(mode.long())
        if self.use_z:
            x = x + self.z_in(self.z_ln(z.float()))
        x = F.gelu(x)                                                                            # [B, d]
        B = x.shape[0]
        h = self.hor.weight[None].expand(B, -1, -1)                                              # [B, n_h, d_h]
        parts = [x[:, None].expand(-1, self.n_h, -1), h]
        if self.d_k:
            assert zk is not None and zk.shape[1] == self.n_h and zk.shape[2] == self.d_k, (None if zk is None else zk.shape)
            parts.append(zk.float())
        return self.net(torch.cat(parts, -1))


def n_params(m):
    return sum(p.numel() for p in m.parameters() if p.requires_grad)


# ----------------------------------------------------------------------------- distributions
def joint_from_marginals(p_hands_off, p_head, p_D):
    """independent product distribution over the 8 states: [...] x3 -> [..., 8] (the pre-registered product baseline)."""
    bits = state_bits(p_hands_off.device, torch.float64)
    ps = torch.stack([p_hands_off.double(), p_head.double(), p_D.double()], -1)                 # [..., 3]
    ps = ps.clamp(0.0, 1.0)
    P = torch.ones(*ps.shape[:-1], N_STATES, dtype=torch.float64, device=ps.device)
    for v in range(3):
        pv = ps[..., v:v + 1]
        P = P * torch.where(bits[:, v] > 0.5, pv, 1.0 - pv)
    return P


def overlap_from_joint(P8, use_hands=True):
    """q = P(U = 1, D = 1); U = hands_off ∨ head_turned (use_hands) or head_turned only (engaged paper-facing).  P8 [..., 8] -> [...]"""
    bits = state_bits(P8.device, P8.dtype)
    U = (bits[:, 0] + bits[:, 1] > 0.5) if use_hands else (bits[:, 1] > 0.5)
    sel = (U & (bits[:, 2] > 0.5)).to(P8.dtype)
    return (P8 * sel).sum(-1)


def marginals_from_joint(P8):
    """-> dict p_hands_off, p_head, p_D, p_U (hands ∨ head), p_U_head (head only)."""
    bits = state_bits(P8.device, P8.dtype)
    out = {}
    for i, v in enumerate(VARS):
        out["p_" + v] = (P8 * bits[:, i]).sum(-1)
    out["p_U"] = (P8 * ((bits[:, 0] + bits[:, 1]) > 0.5).to(P8.dtype)).sum(-1)
    out["p_U_head"] = out["p_head_turned"]
    return out


def J_from_q(q, dt=DT, w=None):
    """J_t = dt * sum_k w_k q_k; q [..., K]."""
    q = q.float()
    if w is None:
        return dt * q.sum(-1)
    return dt * (q * torch.as_tensor(w, dtype=q.dtype, device=q.device)).sum(-1)


# ----------------------------------------------------------------------------- loss
def masked_joint_nll(logits, y3, obs3):
    """logits [B, K, 8]; y3 [B, K, 3] in {0, 1} (any value where unobserved); obs3 [B, K, 3] bool.
    -> nll [B, K] (0 where nothing observed), mask [B, K] (any variable observed)."""
    logp = F.log_softmax(logits.float(), -1)                                                    # [B, K, 8]
    bits = state_bits(logits.device)                                                            # [8, 3]
    y = y3.float()[..., None, :]                                                                 # [B, K, 1, 3]
    o = obs3.bool()[..., None, :]                                                                # [B, K, 1, 3]
    consistent = ((bits[None, None] == y) | ~o).all(-1)                                         # [B, K, 8]
    mask = obs3.bool().any(-1)                                                                   # [B, K]
    lp = torch.where(consistent, logp, torch.full_like(logp, float("-inf")))
    ll = torch.logsumexp(lp, -1)                                                                 # log P(observed configuration)
    nll = torch.where(mask, -ll, torch.zeros_like(ll))
    return nll, mask


def overlap_loss(logits, y3, obs3):
    nll, mask = masked_joint_nll(logits, y3, obs3)
    m = mask.float()
    total = (nll * m).sum() / m.sum().clamp_min(1.0)
    return total, {"joint_nll": total.detach(), "n_obs": m.sum().detach()}


# ----------------------------------------------------------------------------- targets (numpy, strict future: horizon k -> index k-1)
def build_targets(S, P, Pvalid, M, row_ok, n_h=K):
    """S [N, >=K, 2] float (NaN unobserved; hands_on, head_turned), P / Pvalid [N, >=K, 6] bool, M [N, >=K] bool, row_ok [N] bool
    -> y3 [N, K, 3] uint8, obs3 [N, K, 3] bool.  y3[:, k-1] refers to the state at t + k dt (k >= 1) — nothing at or before t."""
    import numpy as np
    S = np.asarray(S, np.float32)[:, :n_h]
    hands_obs = np.isfinite(S[:, :, 0]) & row_ok[:, None] & M[:, :n_h]
    head_obs = np.isfinite(S[:, :, 1]) & row_ok[:, None] & M[:, :n_h]
    Pv = np.asarray(Pvalid, bool)[:, :n_h]
    Pb = np.asarray(P, bool)[:, :n_h] & Pv
    D = Pb.any(2)
    D_obs = Pv.any(2) & M[:, :n_h]
    y3 = np.stack([np.nan_to_num(S[:, :, 0]) < 0.5, np.nan_to_num(S[:, :, 1]) > 0.5, D], -1).astype(np.uint8)
    obs3 = np.stack([hands_obs, head_obs, D_obs], -1)
    return y3, obs3


def realised_overlap(y3, obs3, use_hands=True):
    """y = any_k (U_{t+k} ∧ D_{t+k}) over the K steps; full = every needed variable observed at every step (the pre-registered ranking rows)."""
    import numpy as np
    U = (y3[:, :, 0] > 0) | (y3[:, :, 1] > 0) if use_hands else (y3[:, :, 1] > 0)
    D = y3[:, :, 2] > 0
    need = (obs3[:, :, 0] & obs3[:, :, 1] & obs3[:, :, 2]) if use_hands else (obs3[:, :, 1] & obs3[:, :, 2])
    y = (U & D & need).any(1)
    full = need.all(1)
    return y, full

Xet Storage Details

Size:: 9.14 kB
Xet hash:: 4a102ff98bdee7bb3d6ae17572dc5947ef74378c658c9807327b39ff3cbcaa1f

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""cwm3/models/overlap_head.py — F8 (R2 WP4 batch 2, 2026-09-17; plan 'F8 驾驭员状态 × 道路需求的时间重叠预测'; pre-registration docs/decisions.md
	2026-09-17T19:32:07Z): a DETACHED, horizon-conditioned JOINT head for P(driver unmet state U_{t+k}, road demand D_{t+k} \| H_t) on the frozen ECAN joint model.

	Variables per horizon k = 1..K (0.2 s steps, K = 15 = 3 s), all label-free target states of the existing packs (nothing after t enters the inputs):
	hands_off_{t+k} = S_c[:, k-1, 0] == 0 (results/targets_kin_causal_routes_v2; NaN = unobserved)
	head_turned_{t+k} = S_c[:, k-1, 1] == 1
	D_{t+k} = any road-pressure family (P ∧ Pvalid)[:, k-1] (results/targets_pressure, whitelist CAN margins); observed where any family is valid ∧ M
	The head emits an 8-way categorical over (hands_off, head_turned, D) per horizon: state index s = 4·hands_off + 2·head_turned + D. The four (U, D)
	joint states and q_k = P(U = 1, D = 1) are marginalisations (overlap_from_joint): U = hands_off ∨ head_turned (manual, paper-facing) or U = head_turned
	only (engaged paper-facing: the engaged hand cells are out of the paper and masked from every paper-facing endpoint).
	Inputs: z [B, d_z] = LN([h_u ; z_e ; z_vf]) of the frozen joint model (detached cache), anchor [B, d_a] = label-free current states with validity
	(hands_on0, head0, 6 P0 families + their validity bits), mode [B] regime, optional zk [B, K, d_k] per-horizon features (the `margonly` attribution arm:
	the frozen model's own marginals, WITHOUT trunk features; use_z=False).
	Baseline (pre-registered): the PRODUCT of the frozen model's own marginals — joint_from_marginals(p_hands_off, p_head, p_D) is the independent 8-way
	distribution; overlap_from_joint of it equals p_U · p_D EXACTLY (identity unit test, tests/test_overlap_head.py).
	Loss: masked joint NLL = −log Σ_{s consistent with the observed variables} p_s (the categorical form of a masked BCE: unobserved variables are
	marginalised; a (row, k) with nothing observed contributes 0). J_t = Δt Σ_k w_k q_k (w_k = 1 pre-fixed). ≈ 0.06 M parameters (n_params()).
	No openpilot DMS field anywhere (states = our own pose chains and CAN margins).
	"""
	from __future__ import annotations

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	K = 15 # horizons (0.2 s steps) = 3 s
	DT = 0.2
	VARS = ("hands_off", "head_turned", "D") # bit order of the 8-way state index: s = 4hands_off + 2head_turned + D
	N_STATES = 8
	REPORT_STEPS = {"0.6s": 3, "1s": 5, "2s": 10, "3s": 15} # 0.5 s is not on the grid: step 3 (0.6 s) = nearest step >= 0.5 s
	D_ANCHOR = 2 + 2 + 6 + 6 # hands_on0, head0, their validity, P0 x6, P0valid x6

	# state bits [8, 3]: bit_v(s)
	_BITS = torch.tensor([[(s >> 2) & 1, (s >> 1) & 1, s & 1] for s in range(N_STATES)], dtype=torch.float32)


	def state_bits(device=None, dtype=torch.float32):
	b = _BITS.to(dtype)
	return b.to(device) if device is not None else b


	class OverlapHead(nn.Module):
	def __init__(self, d_z=384, d_anchor=D_ANCHOR, d=96, d_h=16, n_h=K, d_k=0, use_z=True, dropout=0.0):
	super().__init__()
	self.use_z, self.n_h, self.d_k = bool(use_z), int(n_h), int(d_k)
	self.z_ln = nn.LayerNorm(d_z) if use_z else None
	self.z_in = nn.Linear(d_z, d) if use_z else None
	self.a_in = nn.Linear(d_anchor, d)
	self.mode = nn.Embedding(2, d)
	self.hor = nn.Embedding(n_h, d_h)
	self.net = nn.Sequential(nn.Linear(d + d_h + self.d_k, d), nn.GELU(), nn.Dropout(dropout), nn.Linear(d, d), nn.GELU(), nn.Linear(d, N_STATES))

	def forward(self, z, anchor, mode, zk=None):
	"""-> logits [B, n_h, 8]."""
	x = self.a_in(anchor.float()) + self.mode(mode.long())
	if self.use_z:
	x = x + self.z_in(self.z_ln(z.float()))
	x = F.gelu(x) # [B, d]
	B = x.shape[0]
	h = self.hor.weight[None].expand(B, -1, -1) # [B, n_h, d_h]
	parts = [x[:, None].expand(-1, self.n_h, -1), h]
	if self.d_k:
	assert zk is not None and zk.shape[1] == self.n_h and zk.shape[2] == self.d_k, (None if zk is None else zk.shape)
	parts.append(zk.float())
	return self.net(torch.cat(parts, -1))


	def n_params(m):
	return sum(p.numel() for p in m.parameters() if p.requires_grad)


	# ----------------------------------------------------------------------------- distributions
	def joint_from_marginals(p_hands_off, p_head, p_D):
	"""independent product distribution over the 8 states: [...] x3 -> [..., 8] (the pre-registered product baseline)."""
	bits = state_bits(p_hands_off.device, torch.float64)
	ps = torch.stack([p_hands_off.double(), p_head.double(), p_D.double()], -1) # [..., 3]
	ps = ps.clamp(0.0, 1.0)
	P = torch.ones(*ps.shape[:-1], N_STATES, dtype=torch.float64, device=ps.device)
	for v in range(3):
	pv = ps[..., v:v + 1]
	P = P * torch.where(bits[:, v] > 0.5, pv, 1.0 - pv)
	return P


	def overlap_from_joint(P8, use_hands=True):
	"""q = P(U = 1, D = 1); U = hands_off ∨ head_turned (use_hands) or head_turned only (engaged paper-facing). P8 [..., 8] -> [...]"""
	bits = state_bits(P8.device, P8.dtype)
	U = (bits[:, 0] + bits[:, 1] > 0.5) if use_hands else (bits[:, 1] > 0.5)
	sel = (U & (bits[:, 2] > 0.5)).to(P8.dtype)
	return (P8 * sel).sum(-1)


	def marginals_from_joint(P8):
	"""-> dict p_hands_off, p_head, p_D, p_U (hands ∨ head), p_U_head (head only)."""
	bits = state_bits(P8.device, P8.dtype)
	out = {}
	for i, v in enumerate(VARS):
	out["p_" + v] = (P8 * bits[:, i]).sum(-1)
	out["p_U"] = (P8 * ((bits[:, 0] + bits[:, 1]) > 0.5).to(P8.dtype)).sum(-1)
	out["p_U_head"] = out["p_head_turned"]
	return out


	def J_from_q(q, dt=DT, w=None):
	"""J_t = dt * sum_k w_k q_k; q [..., K]."""
	q = q.float()
	if w is None:
	return dt * q.sum(-1)
	return dt * (q * torch.as_tensor(w, dtype=q.dtype, device=q.device)).sum(-1)


	# ----------------------------------------------------------------------------- loss
	def masked_joint_nll(logits, y3, obs3):
	"""logits [B, K, 8]; y3 [B, K, 3] in {0, 1} (any value where unobserved); obs3 [B, K, 3] bool.
	-> nll [B, K] (0 where nothing observed), mask [B, K] (any variable observed)."""
	logp = F.log_softmax(logits.float(), -1) # [B, K, 8]
	bits = state_bits(logits.device) # [8, 3]
	y = y3.float()[..., None, :] # [B, K, 1, 3]
	o = obs3.bool()[..., None, :] # [B, K, 1, 3]
	consistent = ((bits[None, None] == y) \| ~o).all(-1) # [B, K, 8]
	mask = obs3.bool().any(-1) # [B, K]
	lp = torch.where(consistent, logp, torch.full_like(logp, float("-inf")))
	ll = torch.logsumexp(lp, -1) # log P(observed configuration)
	nll = torch.where(mask, -ll, torch.zeros_like(ll))
	return nll, mask


	def overlap_loss(logits, y3, obs3):
	nll, mask = masked_joint_nll(logits, y3, obs3)
	m = mask.float()
	total = (nll * m).sum() / m.sum().clamp_min(1.0)
	return total, {"joint_nll": total.detach(), "n_obs": m.sum().detach()}


	# ----------------------------------------------------------------------------- targets (numpy, strict future: horizon k -> index k-1)
	def build_targets(S, P, Pvalid, M, row_ok, n_h=K):
	"""S [N, >=K, 2] float (NaN unobserved; hands_on, head_turned), P / Pvalid [N, >=K, 6] bool, M [N, >=K] bool, row_ok [N] bool
	-> y3 [N, K, 3] uint8, obs3 [N, K, 3] bool. y3[:, k-1] refers to the state at t + k dt (k >= 1) — nothing at or before t."""
	import numpy as np
	S = np.asarray(S, np.float32)[:, :n_h]
	hands_obs = np.isfinite(S[:, :, 0]) & row_ok[:, None] & M[:, :n_h]
	head_obs = np.isfinite(S[:, :, 1]) & row_ok[:, None] & M[:, :n_h]
	Pv = np.asarray(Pvalid, bool)[:, :n_h]
	Pb = np.asarray(P, bool)[:, :n_h] & Pv
	D = Pb.any(2)
	D_obs = Pv.any(2) & M[:, :n_h]
	y3 = np.stack([np.nan_to_num(S[:, :, 0]) < 0.5, np.nan_to_num(S[:, :, 1]) > 0.5, D], -1).astype(np.uint8)
	obs3 = np.stack([hands_obs, head_obs, D_obs], -1)
	return y3, obs3


	def realised_overlap(y3, obs3, use_hands=True):
	"""y = any_k (U_{t+k} ∧ D_{t+k}) over the K steps; full = every needed variable observed at every step (the pre-registered ranking rows)."""
	import numpy as np
	U = (y3[:, :, 0] > 0) \| (y3[:, :, 1] > 0) if use_hands else (y3[:, :, 1] > 0)
	D = y3[:, :, 2] > 0
	need = (obs3[:, :, 0] & obs3[:, :, 1] & obs3[:, :, 2]) if use_hands else (obs3[:, :, 1] & obs3[:, :, 2])
	y = (U & D & need).any(1)
	full = need.all(1)
	return y, full