Buckets:
| """CWM3-D-kp "K-v2" (plan A.1, 2026-09-14; run family `cwm3d_v2_*`): K's part-token causal trunk kept, plus a joint stream, a lean | |
| shared read-out and an onset-hazard gated anchor mixture. `--model cwm3d` (aide/models/cwm3d.py) is untouched. | |
| obs [B,5,136,3] --jtok--> joint tokens s [B,T,J,d_j] --spatial (per frame, no cross-frame mixing)--> s | |
| obs --tok (heads.PartTokenizer)--> part tokens x [B,T,4,d] (+ pool(s): per-part attention pooling of the joints, residual) | |
| x [B,T*4,d] --layers x heads.Block, causal_part_mask (same frame full, cross-frame causal)--> ln_f --> parts [B,T,4,d]; z = parts[:, -1] | |
| u = jtime(s) [B,J,d_j] (temporal mix of the joint stream: Linear(T*d_j -> d_j) + LN, or LN of the last frame) | |
| head : LeanHorizonHead(u, z) -> [B,H,J,7] (same channel layout as heads.HorizonHead: mu, log_b, 3 anchor logits, gate; zero-init) | |
| or heads.HorizonHead(z) (K's per-part flat MLP, `--no-lean-head`) | |
| haz : HazardHead(z) -> onset-hazard logits [B,4,H] per part; P = 1 - prod_{k<=h}(1 - lam_k) (BATON forward_marginal with h_off = 0) | |
| mod : HazardGate: Pc = P - prior_P; gate' = gate*(1 - beta_h*P) + gamma_h*Pc; mu' = mu*exp(kappa_h*Pc) (beta, gamma, kappa zero-init -> identity) | |
| mix : heads.AnchorMixture (unchanged) -> pred, log_b | |
| Every flag combination is EXACTLY zero-velocity at initialisation (|pred - last| <= 1e-6 px) and a pure function of the 5 observed frames | |
| (aide/tests/test_cwm3d_v2.py). Training-side pieces (hazard targets / BCE / TRAIN prior) live in aide/models/losses.py and aide/train/train_kp.py. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models import heads # noqa: E402 | |
| from aide.models.cwm3d_spatial import SpatialBlock # noqa: E402 | |
| from aide.models.heads import N_CH, AnchorMixture, Block, HorizonHead, PartTokenizer, anchors, causal_part_mask, joint_channels # noqa: E402 | |
| N_PART = len(P.PARTS) | |
| EPS_H = 1e-5 # hazard clamp, == models/heads.py EPS_H | |
| def part_of_joint() -> torch.Tensor: | |
| """[J] long: index of the part (P.PART_NAMES order) each HALPE joint belongs to.""" | |
| poj = torch.empty(P.J, dtype=torch.long) | |
| for p, k in enumerate(P.PART_NAMES): | |
| poj[torch.as_tensor(P.PARTS[k], dtype=torch.long)] = p | |
| return poj | |
| def part_mask() -> torch.Tensor: | |
| """[4,J] bool: True where joint j belongs to part p.""" | |
| return torch.nn.functional.one_hot(part_of_joint(), N_PART).t().bool() | |
| V2_DEFAULTS = dict(d=256, layers=3, heads=4, ffn=1024, d_j=48, s_heads=4, s_blocks=2, joint_time="all", spatial=True, lean_head=True, | |
| hazard=True, haz_hidden=128, haz_mod=True, haz_detach=True, head_hidden=256, d_h=32, k_head_hidden=512) | |
| # ----------------------------------------------------------------------------- joint stream | |
| class JointTokenizer(nn.Module): | |
| """obs [B,T,J,3] -> joint tokens [B,T,J,d_j] = Linear(5 -> d_j)(heads.joint_channels) + joint embedding.""" | |
| def __init__(self, d_j=48): | |
| super().__init__() | |
| self.inp = nn.Linear(N_CH, d_j) | |
| self.joint_emb = nn.Parameter(torch.zeros(P.J, d_j)) | |
| nn.init.normal_(self.joint_emb, std=0.02) | |
| def forward(self, obs): | |
| return self.inp(joint_channels(obs)) + self.joint_emb[None, None] | |
| class SpatialMixer(nn.Module): | |
| """s_blocks x cwm3d_spatial.SpatialBlock over the joint axis of each frame (frames independent), then LN. blocks = 0 -> LN only.""" | |
| def __init__(self, d_j=48, heads=4, blocks=2): | |
| super().__init__() | |
| self.blocks = nn.ModuleList(SpatialBlock(d_j, heads, 4 * d_j) for _ in range(blocks)) | |
| self.ln_s = nn.LayerNorm(d_j) | |
| def forward(self, s): # [B,T,J,d_j] | |
| B, T, J, d = s.shape | |
| x = s.reshape(B * T, J, d) | |
| for blk in self.blocks: | |
| x = blk(x) | |
| return self.ln_s(x).view(B, T, J, d) | |
| class PartAttnPool(nn.Module): | |
| """One learned query per part, masked softmax over that part's joints -> pooled [B,T,4,d_j] -> Linear(d_j -> d): residual for the part tokens.""" | |
| def __init__(self, d_j=48, d=256): | |
| super().__init__() | |
| self.q = nn.Parameter(torch.zeros(N_PART, d_j)) | |
| nn.init.normal_(self.q, std=0.02) | |
| self.proj = nn.Linear(d_j, d) | |
| self.register_buffer("mask", part_mask()) # [4,J] | |
| self.scale = 1.0 / math.sqrt(d_j) | |
| def forward(self, s): # [B,T,J,d_j] -> [B,T,4,d] | |
| att = torch.einsum("pd,btjd->btpj", self.q, s) * self.scale | |
| att = att.masked_fill(~self.mask[None, None], float("-inf")).softmax(-1) | |
| return self.proj(torch.einsum("btpj,btjd->btpd", att, s)) | |
| class JointTimeMix(nn.Module): | |
| """Joint-stream temporal read-out -> u [B,J,d_j]. "all": Linear(T*d_j -> d_j) + LN over the 5 observed frames; "last": LN(s[:, -1]).""" | |
| def __init__(self, d_j=48, mode="all", T=P.T_OBS): | |
| super().__init__() | |
| assert mode in ("all", "last"), mode | |
| self.mode = mode | |
| if mode == "all": | |
| self.mix = nn.Linear(T * d_j, d_j) | |
| self.ln = nn.LayerNorm(d_j) | |
| def forward(self, s): # [B,T,J,d_j] | |
| if self.mode == "last": | |
| return self.ln(s[:, -1]) | |
| B, T, J, d = s.shape | |
| return self.ln(self.mix(s.permute(0, 2, 1, 3).reshape(B, J, T * d))) | |
| # ----------------------------------------------------------------------------- heads | |
| class LeanHorizonHead(nn.Module): | |
| """Shared read-out over (horizon, joint): feature = [u_j | LN(z_last)[part(j)] | h_emb[h] + j_emb[j]] -> LN -> Linear(hidden) -> GELU -> | |
| Linear(hidden -> HorizonHead.N_C). Output layer zero-initialised; output [B,H,J,7] with HorizonHead's channel slices.""" | |
| def __init__(self, d_j=48, d=256, d_h=32, hidden=256, H=P.T_PRED): | |
| super().__init__() | |
| self.H = H | |
| self.ln_z = nn.LayerNorm(d) | |
| self.h_emb = nn.Parameter(torch.zeros(H, d_h)) | |
| self.j_emb = nn.Parameter(torch.zeros(P.J, d_h)) | |
| nn.init.normal_(self.h_emb, std=0.02) | |
| nn.init.normal_(self.j_emb, std=0.02) | |
| d_in = d_j + d + d_h | |
| self.mlp = nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, hidden), nn.GELU(), nn.Linear(hidden, HorizonHead.N_C)) | |
| nn.init.zeros_(self.mlp[-1].weight) | |
| nn.init.zeros_(self.mlp[-1].bias) | |
| self.register_buffer("poj", part_of_joint()) # [J] | |
| def forward(self, u, z): # u [B,J,d_j], z [B,4,d] -> [B,H,J,7] | |
| B = u.shape[0] | |
| zj = self.ln_z(z)[:, self.poj] # [B,J,d] | |
| e = self.h_emb[:, None] + self.j_emb[None] # [H,J,d_h] | |
| f = torch.cat([u[:, None].expand(-1, self.H, -1, -1), zj[:, None].expand(-1, self.H, -1, -1), e[None].expand(B, -1, -1, -1)], -1) | |
| return self.mlp(f) | |
| class HazardHead(nn.Module): | |
| """z_last [B,4,d] -> onset-hazard logits [B,4,H] (shared LN -> Linear -> GELU -> Linear(hidden -> H, zero weight) + per-part bias). | |
| Buffers: tau [4] px, prior_P [4,H] (default 1 - 0.5^h), prior_lam [4,H]. set_prior() writes the TRAIN prior into the buffers and the bias | |
| (= logit of the prior hazard), so the untrained head outputs P == prior_P for every input.""" | |
| def __init__(self, d=256, hidden=128, H=P.T_PRED): | |
| super().__init__() | |
| self.H = H | |
| self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, hidden), nn.GELU(), nn.Linear(hidden, H, bias=False)) | |
| nn.init.zeros_(self.net[-1].weight) | |
| self.bias = nn.Parameter(torch.zeros(N_PART, H)) | |
| self.register_buffer("tau", torch.zeros(N_PART)) | |
| self.register_buffer("prior_P", torch.zeros(N_PART, H)) | |
| self.register_buffer("prior_lam", torch.zeros(N_PART, H)) | |
| self.set_prior(torch.zeros(N_PART), (1 - 0.5 ** torch.arange(1, H + 1, dtype=torch.float32))[None].expand(N_PART, -1)) | |
| def lam_from_P(prior_P): | |
| """cumulative P [.,H] -> per-step hazards lam_h = 1 - (1 - P_h) / (1 - P_{h-1}), clamped to (EPS_H, 1 - EPS_H).""" | |
| surv = (1 - prior_P).clamp_min(EPS_H) | |
| prev = torch.cat([torch.ones_like(surv[..., :1]), surv[..., :-1]], -1) | |
| return (1 - surv / prev).clamp(EPS_H, 1 - EPS_H) | |
| def set_prior(self, tau, prior_P): | |
| tau = torch.as_tensor(tau, dtype=torch.float32).to(self.tau.device) | |
| prior_P = torch.as_tensor(prior_P, dtype=torch.float32).to(self.prior_P.device) | |
| assert tau.shape == (N_PART,) and prior_P.shape == (N_PART, self.H), (tau.shape, prior_P.shape) | |
| self.tau.copy_(tau) | |
| self.prior_P.copy_(prior_P) | |
| self.prior_lam.copy_(self.lam_from_P(prior_P)) | |
| self.bias.copy_(torch.logit(self.prior_lam)) | |
| def forward(self, z): # [B,4,d] -> logits [B,4,H] | |
| return self.net(z) + self.bias[None] | |
| def cumulative(logits): | |
| """P_h = 1 - prod_{k<=h}(1 - lam_k), lam = sigmoid(logits) clamped (== models/heads.py forward_marginal with h_off = 0, z0 = 0).""" | |
| lam = torch.sigmoid(logits).clamp(EPS_H, 1 - EPS_H) | |
| return 1 - torch.cumprod(1 - lam, dim=-1) | |
| class HazardGate(nn.Module): | |
| """Modulates the head output with the cumulative onset probability of the joint's part. Pj [B,H,J]; Pc = Pj - prior_P; | |
| gate' = gate*(1 - beta_h*Pj) + gamma_h*Pc; mu' = mu*exp(kappa_h*Pc); log_b and anchor logits untouched. beta = gamma = kappa = 0 -> identity.""" | |
| def __init__(self, H=P.T_PRED, detach=True): | |
| super().__init__() | |
| self.detach = detach | |
| self.beta = nn.Parameter(torch.zeros(H)) | |
| self.gamma = nn.Parameter(torch.zeros(H)) | |
| self.kappa = nn.Parameter(torch.zeros(H)) | |
| self.register_buffer("poj", part_of_joint()) | |
| def forward(self, ho, P_cum, prior_P): # ho [B,H,J,7], P_cum [B,4,H], prior_P [4,H] | |
| if self.detach: | |
| P_cum = P_cum.detach() | |
| Pj = P_cum[:, self.poj].permute(0, 2, 1) # [B,H,J] | |
| Pc = Pj - prior_P[self.poj].t()[None] # [B,H,J] | |
| gate = ho[..., HorizonHead.C_GATE] * (1 - self.beta[None, :, None] * Pj) + self.gamma[None, :, None] * Pc | |
| mu = ho[..., HorizonHead.C_MU] * torch.exp(self.kappa[None, :, None] * Pc)[..., None] | |
| return torch.cat([mu, ho[..., 2:HorizonHead.C_GATE], gate[..., None]], -1) | |
| # ----------------------------------------------------------------------------- model | |
| class CWM3Dv2(nn.Module): | |
| name = "cwm3d_v2" | |
| # trunk-only initialisation (train_kp --init-parts trunk): everything before the read-out; head / haz / mod keep their seed (zero) init | |
| TRUNK_PREFIXES = ("tok.", "jtok.", "spatial.", "pool.", "jtime.", "blocks.", "ln_f.") | |
| def __init__(self, d=256, layers=3, heads=4, ffn=1024, dropout=0.0, d_j=48, s_heads=4, s_blocks=2, joint_time="all", spatial=True, | |
| lean_head=True, hazard=True, haz_hidden=128, haz_mod=True, haz_detach=True, head_hidden=256, d_h=32, k_head_hidden=512): | |
| super().__init__() | |
| self.d, self.d_j, self.lean_head = d, d_j, lean_head | |
| self.tok = PartTokenizer(d) | |
| self.jtok = JointTokenizer(d_j) | |
| self.spatial = SpatialMixer(d_j, s_heads, s_blocks) if spatial else None | |
| self.pool = PartAttnPool(d_j, d) if spatial else None | |
| self.jtime = JointTimeMix(d_j, joint_time) if lean_head else None | |
| self.blocks = nn.ModuleList(Block(d, heads, ffn, dropout) for _ in range(layers)) | |
| self.ln_f = nn.LayerNorm(d) | |
| self.register_buffer("mask", causal_part_mask(P.T_OBS, N_PART, 0)) | |
| self.head = LeanHorizonHead(d_j, d, d_h, head_hidden, P.T_PRED) if lean_head else HorizonHead(d, k_head_hidden, P.T_PRED) | |
| self.haz = HazardHead(d, haz_hidden, P.T_PRED) if hazard else None | |
| self.mod = HazardGate(P.T_PRED, haz_detach) if (hazard and haz_mod) else None | |
| self.mix = AnchorMixture(enabled=True) | |
| def set_hazard_prior(self, tau, prior_P): | |
| assert self.haz is not None, "no hazard head" | |
| self.haz.set_prior(tau, prior_P) | |
| def encode(self, obs): | |
| """-> part tokens [B,T,4,d] (post-trunk), joint read-out u [B,J,d_j] (None with K's head).""" | |
| B, T = obs.shape[:2] | |
| s = self.jtok(obs) # [B,T,J,d_j] | |
| if self.spatial is not None: | |
| s = self.spatial(s) | |
| x = self.tok(obs) # [B,T,4,d] | |
| if self.pool is not None: | |
| x = x + self.pool(s) | |
| L = T * N_PART | |
| x = x.reshape(B, L, self.d) | |
| mask = self.mask[:L, :L] | |
| for blk in self.blocks: | |
| x = blk(x, mask) | |
| parts = self.ln_f(x).view(B, T, N_PART, self.d) | |
| u = self.jtime(s) if self.jtime is not None else None | |
| return parts, u | |
| def forward(self, obs): | |
| assert obs.shape[1] == P.T_OBS and obs.shape[2] == P.J, obs.shape | |
| out = {"last": obs[:, -1:, :, :2]} | |
| parts, u = self.encode(obs) | |
| z = parts[:, -1] # [B,4,d] | |
| ho = self.head(u, z) if self.lean_head else self.head(z) # [B,H,J,7] | |
| if self.haz is not None: | |
| lam = self.haz(z) # [B,4,H] onset-hazard logits | |
| P_cum = self.haz.cumulative(lam) | |
| out["haz_logits"], out["haz_P"] = lam, P_cum | |
| if self.mod is not None: | |
| ho = self.mod(ho, P_cum, self.haz.prior_P) | |
| pred, log_b = self.mix(ho, anchors(obs)) | |
| out["pred"], out["log_b"] = pred, log_b | |
| out["anchor_w"] = torch.softmax(ho[..., HorizonHead.C_LOGIT], -1) | |
| out["anchor_gate"] = ho[..., HorizonHead.C_GATE] | |
| return out | |
| # ----------------------------------------------------------------------------- CLI plumbing (train_kp.py / pretrain_cwm3d_baton.py) | |
| def add_args(ap): | |
| """K-v2 flags. Re-uses the existing --d / --heads / --ffn / --head-hidden (K's head width) of the caller; --layers is NOT re-used.""" | |
| g = ap.add_argument_group("cwm3d_v2") | |
| g.add_argument("--v2-layers", type=int, default=V2_DEFAULTS["layers"], help="part-trunk depth (K uses --layers 4)") | |
| g.add_argument("--v2-dj", type=int, default=V2_DEFAULTS["d_j"], help="joint-stream width") | |
| g.add_argument("--v2-sheads", type=int, default=V2_DEFAULTS["s_heads"]) | |
| g.add_argument("--v2-sblocks", type=int, default=V2_DEFAULTS["s_blocks"], help="spatial (joint-axis) attention blocks per frame") | |
| g.add_argument("--v2-joint-time", default=V2_DEFAULTS["joint_time"], choices=("all", "last"), help="joint-stream temporal read-out") | |
| g.add_argument("--v2-head-hidden", type=int, default=V2_DEFAULTS["head_hidden"], help="lean head hidden width") | |
| g.add_argument("--v2-dh", type=int, default=V2_DEFAULTS["d_h"], help="horizon / joint embedding width of the lean head") | |
| g.add_argument("--no-v2-spatial", action="store_true", help="drop the spatial mixer and the part attention pooling") | |
| g.add_argument("--no-lean-head", action="store_true", help="K's per-part flat HorizonHead (width --head-hidden) instead of the lean read-out") | |
| g.add_argument("--no-hazard", action="store_true", help="no onset-hazard head (and no modulation)") | |
| g.add_argument("--haz-weight", type=float, default=1.0, help="weight of the onset-hazard BCE in the training loss (0 = off)") | |
| g.add_argument("--haz-quantile", type=float, default=0.75, help="TRAIN quantile of the h=5 part-centroid displacement defining tau_p") | |
| g.add_argument("--haz-hidden", type=int, default=V2_DEFAULTS["haz_hidden"]) | |
| g.add_argument("--haz-mod", type=int, default=1, choices=(0, 1), help="1: hazard-gated anchor mixture; 0: hazard head as auxiliary loss only") | |
| g.add_argument("--haz-detach", type=int, default=1, choices=(0, 1), help="1: P enters the modulation as a constant (geometry loss does not train the hazard head)") | |
| return ap | |
| def kwargs_from_args(a) -> dict: | |
| """argparse namespace -> CWM3Dv2 kwargs (missing attributes fall back to V2_DEFAULTS; dropout defaults to 0).""" | |
| g = lambda k, dflt: getattr(a, k, dflt) # noqa: E731 | |
| D = V2_DEFAULTS | |
| return dict(d=g("d", D["d"]), layers=g("v2_layers", D["layers"]), heads=g("heads", D["heads"]), ffn=g("ffn", D["ffn"]), | |
| dropout=float(g("dropout", 0.0)), d_j=g("v2_dj", D["d_j"]), s_heads=g("v2_sheads", D["s_heads"]), s_blocks=g("v2_sblocks", D["s_blocks"]), | |
| joint_time=g("v2_joint_time", D["joint_time"]), spatial=not g("no_v2_spatial", False), lean_head=not g("no_lean_head", False), | |
| hazard=not g("no_hazard", False), haz_hidden=g("haz_hidden", D["haz_hidden"]), haz_mod=bool(g("haz_mod", 1)), | |
| haz_detach=bool(g("haz_detach", 1)), head_hidden=g("v2_head_hidden", D["head_hidden"]), d_h=g("v2_dh", D["d_h"]), | |
| k_head_hidden=g("head_hidden", D["k_head_hidden"])) | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "cwm3d_v2": | |
| return CWM3Dv2(**kw) | |
| raise ValueError(name) | |
| def param_breakdown(m: nn.Module) -> dict: | |
| out = {} | |
| for n, mod in m.named_children(): | |
| k = sum(p.numel() for p in mod.parameters()) | |
| if k: | |
| out[n] = k | |
| out["total"] = sum(p.numel() for p in m.parameters()) | |
| return out | |
| if __name__ == "__main__": | |
| for kw in ({}, {"spatial": False}, {"lean_head": False}, {"hazard": False}): | |
| m = CWM3Dv2(**kw) | |
| x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| o = m(x) | |
| print(kw or "full", param_breakdown(m), "zero-init == ZV:", (o["pred"] - x[:, -1:, :, :2]).abs().max().item() <= 1e-6) | |
Xet Storage Details
- Size:
- 17.9 kB
- Xet hash:
- 741382417ba876e730009773c80b90396b3492a362a0860b829ec3b29cc39a99
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.