Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /models /kp_forecaster.py

AsianPlayer

16 days ago

download

raw

13.6 kB

"""cwm3/models/kp_forecaster.py — LANE F1 keypoint RESIDUAL forecaster on BATON (R2 WP4 batch 2; pre-registration docs/decisions.md 2026-09-17T19:33:12Z).
Observation (all frames <= t, 10 Hz ticks relative to the anchor):  `full`  = dense last 2 s (ticks -19..0, 20 frames) + sparse 1 Hz history (ticks -30, -40, …,
-100, 8 frames);  `dense` = the 20 dense frames only;  `1hz` = ticks 0, -10, …, -100 (11 frames = the observation DOWNSAMPLED to 1 Hz).  Every frame carries the
11 upper-body joints (x/W, y/H, conf, valid) and its time gap Δt = t_frame − t (<= 0); the H = 4 horizon QUERY frames (+0.5 / +1 / +2 / +3 s) carry zeros and
Δt = +h.  CAUSAL NORMALISATION (`causal_normalise`): per joint subtract the LAST VALID position <= t (`ref`), divide by the median shoulder width over the observed
frames (`scale`, fallback constant when no frame has both shoulders; clamped) — no future frame enters (tests/test_kp_forecaster.py poisons the future).
Output: residual of the future joints relative to `ref` in normalised units, M = 3 trajectory MODES [B,H,M,J,2] + mode logits [B,M] + state logits [B,3,2]
(P(hands_on), P(head_turned) at +1 / +2 / +3 s) — the distillation channel (state probabilities / modes, never the mean coordinates).
Backbones (parameter-matched ±10 %, both random init):
  `mb`  = the MotionBERT ARCHITECTURE: the official DSTformer modules (aide/models/motionbert_official/DSTformer.py — dual-stream spatial / temporal blocks with
          the att_fuse mixing, joints_embed + pos_embed + temp_embed) in a small config, plus a time-gap Linear on Δt; NOT the released weights.
  `cwm` = the CockpitWM part-tokenised causal temporal-part Transformer form: 4 part tokens per frame (head, torso, left arm, right arm), causal mask over frames,
          horizon read-out from the query frames' part tokens.
Loss (`kp_loss`): winner-takes-all masked L1 over the modes + CE on the arg-min mode + masked BCE on the future states.  No CAN column, label or DMS field is read here.
"""
from __future__ import annotations

import sys
from pathlib import Path

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

CWM3 = Path(__file__).resolve().parents[1]
if str(CWM3) not in sys.path:
    sys.path.insert(0, str(CWM3))
from aide.models.motionbert_official.DSTformer import DSTformer   # noqa: E402  (architecture only; no weights)

NJ = 11
N_IN = 5                                                      # xn_x, xn_y, conf, valid, dt/10
HORIZON_TICKS = (5, 10, 20, 30)                                # +0.5 / +1 / +2 / +3 s at 10 Hz
HORIZONS_S = tuple(k / 10.0 for k in HORIZON_TICKS)
H = len(HORIZON_TICKS)
N_MODES = 3
N_STATES, N_STATE_H = 2, 3                                     # hands_on / head_turned at +1 / +2 / +3 s
L_SHO, R_SHO = 5, 6
SCALE_FALLBACK, SCALE_MIN, SCALE_MAX = 0.12, 0.03, 0.50
PARTS = {"head": [0, 1, 2, 3, 4], "torso": [5, 6], "larm": [5, 7, 9], "rarm": [6, 8, 10]}
PART_OF_JOINT = [0, 0, 0, 0, 0, 1, 1, 2, 3, 2, 3]             # read-out part per joint (shoulders from torso)
KPF_DIM = N_STATES * N_STATE_H + N_MODES + 1                   # distillation channel: 6 state probs + 3 mode probs + valid bit = 10


def obs_offsets(kind):
    """relative 10 Hz ticks of the observed frames (ascending, all <= 0)."""
    if kind == "full":
        off = list(range(-100, -20, 10)) + list(range(-19, 1))
    elif kind == "dense":
        off = list(range(-19, 1))
    elif kind == "1hz":
        off = list(range(-100, 1, 10))
    else:
        raise ValueError(kind)
    assert max(off) == 0 and all(o <= 0 for o in off)
    return torch.tensor(off, dtype=torch.long)


def causal_normalise(xy, valid, scale_fallback=SCALE_FALLBACK):
    """xy [B,T,J,2] (frames ascending in time, the last one = t), valid [B,T,J] bool -> xn [B,T,J,2], ref [B,J,2], ref_ok [B,J], scale [B].
    ref = last valid position <= t per joint (0 when never valid); scale = median shoulder width over frames with both shoulders valid (fallback
    constant), clamped to [SCALE_MIN, SCALE_MAX].  Only the given (observed) frames are read."""
    B, T, J, _ = xy.shape
    xy = torch.nan_to_num(xy.float())
    ref = torch.zeros(B, J, 2, device=xy.device)
    ref_ok = torch.zeros(B, J, dtype=torch.bool, device=xy.device)
    for k in range(T):                                        # oldest -> newest: the newest valid frame wins
        v = valid[:, k]
        ref = torch.where(v[..., None], xy[:, k], ref)
        ref_ok = ref_ok | v
    sw = (xy[:, :, L_SHO] - xy[:, :, R_SHO]).norm(dim=-1)                       # [B,T]
    both = valid[:, :, L_SHO] & valid[:, :, R_SHO]
    sw = torch.where(both, sw, torch.full_like(sw, float("nan")))
    med = torch.nanmedian(sw, dim=1).values
    scale = torch.where(torch.isfinite(med), med, torch.full_like(med, float(scale_fallback))).clamp(SCALE_MIN, SCALE_MAX)
    xn = (xy - ref[:, None]) / scale[:, None, None, None] * valid[..., None].float()
    return xn, ref, ref_ok, scale


def make_tokens(xn, conf, valid, dt_s):
    """-> [B,T,J,N_IN] input channels for the observed frames.  dt_s [T] (<= 0)."""
    B, T, J, _ = xn.shape
    dt = dt_s.to(xn.device).float()[None, :, None, None].expand(B, T, J, 1) / 10.0
    return torch.cat([xn, torch.nan_to_num(conf.float())[..., None] * valid[..., None].float(), valid[..., None].float(), dt], -1)


def query_tokens(B, J, device):
    dt = torch.tensor(HORIZONS_S, device=device).float()[None, :, None, None].expand(B, H, J, 1) / 10.0
    return torch.cat([torch.zeros(B, H, J, N_IN - 1, device=device), dt], -1)


class MBBackbone(nn.Module):
    """MotionBERT-architecture (official DSTformer modules) with a time-gap embedding; returns per-token features [B,F,J,d]."""

    def __init__(self, d=128, depth=3, heads=4, mlp_ratio=2, maxlen=64, dropout=0.0):
        super().__init__()
        self.net = DSTformer(dim_in=N_IN, dim_out=0, dim_feat=d, dim_rep=0, depth=depth, num_heads=heads, mlp_ratio=mlp_ratio, num_joints=NJ, maxlen=maxlen,
                             drop_rate=dropout, attn_drop_rate=dropout, drop_path_rate=0.0, att_fuse=True)
        self.gap = nn.Linear(1, d)                             # time-gap embedding on Δt (seconds / 10), shared by observed and query frames
        self.d = d

    def forward(self, x, dt):                                  # x [B,F,J,N_IN], dt [B,F] (seconds)
        n = self.net
        B, Fr, J, C = x.shape
        h = n.joints_embed(x.reshape(-1, J, C)) + n.pos_embed
        h = h.reshape(B, Fr, J, -1) + n.temp_embed[:, :Fr] + self.gap((dt / 10.0)[..., None])[:, :, None, :]
        h = n.pos_drop(h.reshape(B * Fr, J, -1))
        for idx, (blk_st, blk_ts) in enumerate(zip(n.blocks_st, n.blocks_ts)):
            x_st, x_ts = blk_st(h, Fr), blk_ts(h, Fr)
            alpha = n.ts_attn[idx](torch.cat([x_st, x_ts], -1)).softmax(-1)
            h = x_st * alpha[:, :, 0:1] + x_ts * alpha[:, :, 1:2]
        return n.norm(h).reshape(B, Fr, J, -1)


class TBlock(nn.Module):
    def __init__(self, d, heads, ffn, dropout=0.0):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, dropout=dropout, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(d, ffn), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn, d))
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.drop(self.attn(h, h, h, attn_mask=mask, need_weights=False)[0])
        return x + self.drop(self.ff(self.ln2(x)))


def causal_part_mask(T, n_part):
    frame = torch.arange(T * n_part) // n_part
    return frame[None, :] > frame[:, None]                     # True = blocked (later frames)


class CWMBackbone(nn.Module):
    """CockpitWM form: part tokens per frame -> causal temporal-part Transformer -> per-joint features read from the joint's part token [B,F,J,d]."""

    def __init__(self, d=128, layers=3, heads=4, ffn=256, maxlen=64, dropout=0.0):
        super().__init__()
        self.parts = list(PARTS.values())
        self.tok = nn.ModuleList([nn.Linear(len(p) * N_IN, d) for p in self.parts])
        self.part_emb = nn.Embedding(len(self.parts), d)
        self.slot_emb = nn.Embedding(maxlen, d)
        self.gap = nn.Linear(1, d)
        self.blocks = nn.ModuleList([TBlock(d, heads, ffn, dropout) for _ in range(layers)])
        self.norm = nn.LayerNorm(d)
        self.joint_emb = nn.Embedding(NJ, d)                   # distinguishes the joints read from one part token
        self.register_buffer("part_of_joint", torch.tensor(PART_OF_JOINT), persistent=False)
        self.d = d

    def forward(self, x, dt):                                  # x [B,F,J,N_IN], dt [B,F]
        B, Fr, J, C = x.shape
        toks = torch.stack([lin(x[:, :, p].reshape(B, Fr, -1)) for lin, p in zip(self.tok, self.parts)], 2)   # [B,F,P,d]
        P = toks.shape[2]
        toks = toks + self.part_emb.weight[None, None] + self.slot_emb.weight[None, :Fr, None] + self.gap((dt / 10.0)[..., None])[:, :, None, :]
        h = toks.reshape(B, Fr * P, -1)
        mask = causal_part_mask(Fr, P).to(h.device)
        for blk in self.blocks:
            h = blk(h, mask)
        h = self.norm(h).reshape(B, Fr, P, -1)
        return h[:, :, self.part_of_joint] + self.joint_emb.weight[None, None]


class KpForecaster(nn.Module):
    """observation tokens + query tokens -> {res [B,H,M,J,2], mode_logits [B,M], state_logits [B,3,2]}."""

    def __init__(self, backbone="mb", obs="full", d=128, depth=3, heads=4, mlp_ratio=2, ffn=256, dropout=0.0, scale_fallback=SCALE_FALLBACK):
        super().__init__()
        self.obs_kind, self.scale_fallback = obs, float(scale_fallback)
        self.register_buffer("offsets", obs_offsets(obs), persistent=False)
        self.n_obs = int(self.offsets.numel())
        if backbone == "mb":
            self.backbone = MBBackbone(d, depth, heads, mlp_ratio, maxlen=self.n_obs + H, dropout=dropout)
        elif backbone == "cwm":
            self.backbone = CWMBackbone(d, depth, heads, ffn, maxlen=self.n_obs + H, dropout=dropout)
        else:
            raise ValueError(backbone)
        self.backbone_kind = backbone
        self.res_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, N_MODES * 2))
        nn.init.zeros_(self.res_head[1].weight)                # zero-velocity at init (every mode = the current pose)
        nn.init.zeros_(self.res_head[1].bias)
        self.pool_head = nn.Sequential(nn.LayerNorm(d * H), nn.Linear(d * H, 64), nn.GELU(), nn.Linear(64, N_MODES + N_STATES * N_STATE_H))

    def forward(self, xy, conf, valid):
        """xy [B,T,J,2] (T = n_obs observed frames, ascending), conf [B,T,J], valid [B,T,J] -> dict (+ ref / ref_ok / scale of the normalisation)."""
        B, T, J, _ = xy.shape
        assert T == self.n_obs and J == NJ, (xy.shape, self.n_obs)
        xn, ref, ref_ok, scale = causal_normalise(xy, valid, self.scale_fallback)
        dt_obs = self.offsets.float() / 10.0                   # seconds, <= 0
        x = torch.cat([make_tokens(xn, conf, valid, dt_obs), query_tokens(B, J, xy.device)], 1)          # [B,T+H,J,N_IN]
        dt = torch.cat([dt_obs.to(xy.device)[None].expand(B, -1), torch.tensor(HORIZONS_S, device=xy.device)[None].expand(B, -1)], 1)
        feat = self.backbone(x, dt)                            # [B,T+H,J,d]
        q = feat[:, T:]                                        # query frames [B,H,J,d]
        res = self.res_head(q).view(B, H, J, N_MODES, 2).permute(0, 1, 3, 2, 4).contiguous()             # [B,H,M,J,2]
        pooled = self.pool_head(q.mean(2).reshape(B, -1))
        return {"res": res, "mode_logits": pooled[:, :N_MODES], "state_logits": pooled[:, N_MODES:].view(B, N_STATE_H, N_STATES),
                "ref": ref, "ref_ok": ref_ok, "scale": scale}


def kp_loss(out, fut_xy, fut_valid, S, lam_ce=0.1, lam_state=0.5):
    """fut_xy [B,H,J,2] frame units, fut_valid [B,H,J], S [B,3,2] future states (NaN censored) -> total, parts, per-mode L1 [B,M]."""
    ref, ref_ok, scale = out["ref"], out["ref_ok"], out["scale"]
    tgt = (torch.nan_to_num(fut_xy.float()) - ref[:, None]) / scale[:, None, None, None]               # normalised residual target
    m = (fut_valid & ref_ok[:, None]).float()[:, :, None, :, None]                                      # [B,H,1,J,1]
    err = ((out["res"] - tgt[:, :, None]).abs() * m).sum((1, 3, 4)) / m.sum((1, 3, 4)).clamp_min(1.0)   # [B,M]
    has = (m.sum((1, 2, 3, 4)) > 0)
    best = err.detach().argmin(1)
    wta = (err.gather(1, best[:, None]).squeeze(1) * has.float()).sum() / has.float().sum().clamp_min(1.0)
    ce = F.cross_entropy(out["mode_logits"][has], best[has]) if has.any() else out["mode_logits"].sum() * 0
    sm = torch.isfinite(S)
    sbce = (F.binary_cross_entropy_with_logits(out["state_logits"], torch.nan_to_num(S), reduction="none") * sm.float()).sum() / sm.float().sum().clamp_min(1.0)
    total = wta + lam_ce * ce + lam_state * sbce
    return total, {"wta_l1": float(wta.detach()), "mode_ce": float(ce.detach()), "state_bce": float(sbce.detach())}, err


def kpf_channel(out, anchor_ok):
    """the distillation channel [B,KPF_DIM]: 6 state probabilities, 3 mode probabilities, valid bit (0 -> every other entry 0)."""
    p_state = torch.sigmoid(out["state_logits"]).reshape(-1, N_STATES * N_STATE_H)
    p_mode = torch.softmax(out["mode_logits"], -1)
    ok = anchor_ok.float()[:, None]
    return torch.cat([p_state * ok, p_mode * ok, ok], 1)


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

Xet Storage Details

Size:: 13.6 kB
Xet hash:: da43b733d37a14d91aeb1e38f6bfa1567d517a6facdfcfe60a073e54299f34ad

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

	"""cwm3/models/kp_forecaster.py — LANE F1 keypoint RESIDUAL forecaster on BATON (R2 WP4 batch 2; pre-registration docs/decisions.md 2026-09-17T19:33:12Z).

	Observation (all frames <= t, 10 Hz ticks relative to the anchor): `full` = dense last 2 s (ticks -19..0, 20 frames) + sparse 1 Hz history (ticks -30, -40, …,
	-100, 8 frames); `dense` = the 20 dense frames only; `1hz` = ticks 0, -10, …, -100 (11 frames = the observation DOWNSAMPLED to 1 Hz). Every frame carries the
	11 upper-body joints (x/W, y/H, conf, valid) and its time gap Δt = t_frame − t (<= 0); the H = 4 horizon QUERY frames (+0.5 / +1 / +2 / +3 s) carry zeros and
	Δt = +h. CAUSAL NORMALISATION (`causal_normalise`): per joint subtract the LAST VALID position <= t (`ref`), divide by the median shoulder width over the observed
	frames (`scale`, fallback constant when no frame has both shoulders; clamped) — no future frame enters (tests/test_kp_forecaster.py poisons the future).
	Output: residual of the future joints relative to `ref` in normalised units, M = 3 trajectory MODES [B,H,M,J,2] + mode logits [B,M] + state logits [B,3,2]
	(P(hands_on), P(head_turned) at +1 / +2 / +3 s) — the distillation channel (state probabilities / modes, never the mean coordinates).
	Backbones (parameter-matched ±10 %, both random init):
	`mb` = the MotionBERT ARCHITECTURE: the official DSTformer modules (aide/models/motionbert_official/DSTformer.py — dual-stream spatial / temporal blocks with
	the att_fuse mixing, joints_embed + pos_embed + temp_embed) in a small config, plus a time-gap Linear on Δt; NOT the released weights.
	`cwm` = the CockpitWM part-tokenised causal temporal-part Transformer form: 4 part tokens per frame (head, torso, left arm, right arm), causal mask over frames,
	horizon read-out from the query frames' part tokens.
	Loss (`kp_loss`): winner-takes-all masked L1 over the modes + CE on the arg-min mode + masked BCE on the future states. No CAN column, label or DMS field is read here.
	"""
	from __future__ import annotations

	import sys
	from pathlib import Path

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

	CWM3 = Path(__file__).resolve().parents[1]
	if str(CWM3) not in sys.path:
	sys.path.insert(0, str(CWM3))
	from aide.models.motionbert_official.DSTformer import DSTformer # noqa: E402 (architecture only; no weights)

	NJ = 11
	N_IN = 5 # xn_x, xn_y, conf, valid, dt/10
	HORIZON_TICKS = (5, 10, 20, 30) # +0.5 / +1 / +2 / +3 s at 10 Hz
	HORIZONS_S = tuple(k / 10.0 for k in HORIZON_TICKS)
	H = len(HORIZON_TICKS)
	N_MODES = 3
	N_STATES, N_STATE_H = 2, 3 # hands_on / head_turned at +1 / +2 / +3 s
	L_SHO, R_SHO = 5, 6
	SCALE_FALLBACK, SCALE_MIN, SCALE_MAX = 0.12, 0.03, 0.50
	PARTS = {"head": [0, 1, 2, 3, 4], "torso": [5, 6], "larm": [5, 7, 9], "rarm": [6, 8, 10]}
	PART_OF_JOINT = [0, 0, 0, 0, 0, 1, 1, 2, 3, 2, 3] # read-out part per joint (shoulders from torso)
	KPF_DIM = N_STATES * N_STATE_H + N_MODES + 1 # distillation channel: 6 state probs + 3 mode probs + valid bit = 10


	def obs_offsets(kind):
	"""relative 10 Hz ticks of the observed frames (ascending, all <= 0)."""
	if kind == "full":
	off = list(range(-100, -20, 10)) + list(range(-19, 1))
	elif kind == "dense":
	off = list(range(-19, 1))
	elif kind == "1hz":
	off = list(range(-100, 1, 10))
	else:
	raise ValueError(kind)
	assert max(off) == 0 and all(o <= 0 for o in off)
	return torch.tensor(off, dtype=torch.long)


	def causal_normalise(xy, valid, scale_fallback=SCALE_FALLBACK):
	"""xy [B,T,J,2] (frames ascending in time, the last one = t), valid [B,T,J] bool -> xn [B,T,J,2], ref [B,J,2], ref_ok [B,J], scale [B].
	ref = last valid position <= t per joint (0 when never valid); scale = median shoulder width over frames with both shoulders valid (fallback
	constant), clamped to [SCALE_MIN, SCALE_MAX]. Only the given (observed) frames are read."""
	B, T, J, _ = xy.shape
	xy = torch.nan_to_num(xy.float())
	ref = torch.zeros(B, J, 2, device=xy.device)
	ref_ok = torch.zeros(B, J, dtype=torch.bool, device=xy.device)
	for k in range(T): # oldest -> newest: the newest valid frame wins
	v = valid[:, k]
	ref = torch.where(v[..., None], xy[:, k], ref)
	ref_ok = ref_ok \| v
	sw = (xy[:, :, L_SHO] - xy[:, :, R_SHO]).norm(dim=-1) # [B,T]
	both = valid[:, :, L_SHO] & valid[:, :, R_SHO]
	sw = torch.where(both, sw, torch.full_like(sw, float("nan")))
	med = torch.nanmedian(sw, dim=1).values
	scale = torch.where(torch.isfinite(med), med, torch.full_like(med, float(scale_fallback))).clamp(SCALE_MIN, SCALE_MAX)
	xn = (xy - ref[:, None]) / scale[:, None, None, None] * valid[..., None].float()
	return xn, ref, ref_ok, scale


	def make_tokens(xn, conf, valid, dt_s):
	"""-> [B,T,J,N_IN] input channels for the observed frames. dt_s [T] (<= 0)."""
	B, T, J, _ = xn.shape
	dt = dt_s.to(xn.device).float()[None, :, None, None].expand(B, T, J, 1) / 10.0
	return torch.cat([xn, torch.nan_to_num(conf.float())[..., None] * valid[..., None].float(), valid[..., None].float(), dt], -1)


	def query_tokens(B, J, device):
	dt = torch.tensor(HORIZONS_S, device=device).float()[None, :, None, None].expand(B, H, J, 1) / 10.0
	return torch.cat([torch.zeros(B, H, J, N_IN - 1, device=device), dt], -1)


	class MBBackbone(nn.Module):
	"""MotionBERT-architecture (official DSTformer modules) with a time-gap embedding; returns per-token features [B,F,J,d]."""

	def __init__(self, d=128, depth=3, heads=4, mlp_ratio=2, maxlen=64, dropout=0.0):
	super().__init__()
	self.net = DSTformer(dim_in=N_IN, dim_out=0, dim_feat=d, dim_rep=0, depth=depth, num_heads=heads, mlp_ratio=mlp_ratio, num_joints=NJ, maxlen=maxlen,
	drop_rate=dropout, attn_drop_rate=dropout, drop_path_rate=0.0, att_fuse=True)
	self.gap = nn.Linear(1, d) # time-gap embedding on Δt (seconds / 10), shared by observed and query frames
	self.d = d

	def forward(self, x, dt): # x [B,F,J,N_IN], dt [B,F] (seconds)
	n = self.net
	B, Fr, J, C = x.shape
	h = n.joints_embed(x.reshape(-1, J, C)) + n.pos_embed
	h = h.reshape(B, Fr, J, -1) + n.temp_embed[:, :Fr] + self.gap((dt / 10.0)[..., None])[:, :, None, :]
	h = n.pos_drop(h.reshape(B * Fr, J, -1))
	for idx, (blk_st, blk_ts) in enumerate(zip(n.blocks_st, n.blocks_ts)):
	x_st, x_ts = blk_st(h, Fr), blk_ts(h, Fr)
	alpha = n.ts_attn[idx](torch.cat([x_st, x_ts], -1)).softmax(-1)
	h = x_st * alpha[:, :, 0:1] + x_ts * alpha[:, :, 1:2]
	return n.norm(h).reshape(B, Fr, J, -1)


	class TBlock(nn.Module):
	def __init__(self, d, heads, ffn, dropout=0.0):
	super().__init__()
	self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
	self.attn = nn.MultiheadAttention(d, heads, dropout=dropout, batch_first=True)
	self.ff = nn.Sequential(nn.Linear(d, ffn), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn, d))
	self.drop = nn.Dropout(dropout)

	def forward(self, x, mask):
	h = self.ln1(x)
	x = x + self.drop(self.attn(h, h, h, attn_mask=mask, need_weights=False)[0])
	return x + self.drop(self.ff(self.ln2(x)))


	def causal_part_mask(T, n_part):
	frame = torch.arange(T * n_part) // n_part
	return frame[None, :] > frame[:, None] # True = blocked (later frames)


	class CWMBackbone(nn.Module):
	"""CockpitWM form: part tokens per frame -> causal temporal-part Transformer -> per-joint features read from the joint's part token [B,F,J,d]."""

	def __init__(self, d=128, layers=3, heads=4, ffn=256, maxlen=64, dropout=0.0):
	super().__init__()
	self.parts = list(PARTS.values())
	self.tok = nn.ModuleList([nn.Linear(len(p) * N_IN, d) for p in self.parts])
	self.part_emb = nn.Embedding(len(self.parts), d)
	self.slot_emb = nn.Embedding(maxlen, d)
	self.gap = nn.Linear(1, d)
	self.blocks = nn.ModuleList([TBlock(d, heads, ffn, dropout) for _ in range(layers)])
	self.norm = nn.LayerNorm(d)
	self.joint_emb = nn.Embedding(NJ, d) # distinguishes the joints read from one part token
	self.register_buffer("part_of_joint", torch.tensor(PART_OF_JOINT), persistent=False)
	self.d = d

	def forward(self, x, dt): # x [B,F,J,N_IN], dt [B,F]
	B, Fr, J, C = x.shape
	toks = torch.stack([lin(x[:, :, p].reshape(B, Fr, -1)) for lin, p in zip(self.tok, self.parts)], 2) # [B,F,P,d]
	P = toks.shape[2]
	toks = toks + self.part_emb.weight[None, None] + self.slot_emb.weight[None, :Fr, None] + self.gap((dt / 10.0)[..., None])[:, :, None, :]
	h = toks.reshape(B, Fr * P, -1)
	mask = causal_part_mask(Fr, P).to(h.device)
	for blk in self.blocks:
	h = blk(h, mask)
	h = self.norm(h).reshape(B, Fr, P, -1)
	return h[:, :, self.part_of_joint] + self.joint_emb.weight[None, None]


	class KpForecaster(nn.Module):
	"""observation tokens + query tokens -> {res [B,H,M,J,2], mode_logits [B,M], state_logits [B,3,2]}."""

	def __init__(self, backbone="mb", obs="full", d=128, depth=3, heads=4, mlp_ratio=2, ffn=256, dropout=0.0, scale_fallback=SCALE_FALLBACK):
	super().__init__()
	self.obs_kind, self.scale_fallback = obs, float(scale_fallback)
	self.register_buffer("offsets", obs_offsets(obs), persistent=False)
	self.n_obs = int(self.offsets.numel())
	if backbone == "mb":
	self.backbone = MBBackbone(d, depth, heads, mlp_ratio, maxlen=self.n_obs + H, dropout=dropout)
	elif backbone == "cwm":
	self.backbone = CWMBackbone(d, depth, heads, ffn, maxlen=self.n_obs + H, dropout=dropout)
	else:
	raise ValueError(backbone)
	self.backbone_kind = backbone
	self.res_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, N_MODES * 2))
	nn.init.zeros_(self.res_head[1].weight) # zero-velocity at init (every mode = the current pose)
	nn.init.zeros_(self.res_head[1].bias)
	self.pool_head = nn.Sequential(nn.LayerNorm(d * H), nn.Linear(d * H, 64), nn.GELU(), nn.Linear(64, N_MODES + N_STATES * N_STATE_H))

	def forward(self, xy, conf, valid):
	"""xy [B,T,J,2] (T = n_obs observed frames, ascending), conf [B,T,J], valid [B,T,J] -> dict (+ ref / ref_ok / scale of the normalisation)."""
	B, T, J, _ = xy.shape
	assert T == self.n_obs and J == NJ, (xy.shape, self.n_obs)
	xn, ref, ref_ok, scale = causal_normalise(xy, valid, self.scale_fallback)
	dt_obs = self.offsets.float() / 10.0 # seconds, <= 0
	x = torch.cat([make_tokens(xn, conf, valid, dt_obs), query_tokens(B, J, xy.device)], 1) # [B,T+H,J,N_IN]
	dt = torch.cat([dt_obs.to(xy.device)[None].expand(B, -1), torch.tensor(HORIZONS_S, device=xy.device)[None].expand(B, -1)], 1)
	feat = self.backbone(x, dt) # [B,T+H,J,d]
	q = feat[:, T:] # query frames [B,H,J,d]
	res = self.res_head(q).view(B, H, J, N_MODES, 2).permute(0, 1, 3, 2, 4).contiguous() # [B,H,M,J,2]
	pooled = self.pool_head(q.mean(2).reshape(B, -1))
	return {"res": res, "mode_logits": pooled[:, :N_MODES], "state_logits": pooled[:, N_MODES:].view(B, N_STATE_H, N_STATES),
	"ref": ref, "ref_ok": ref_ok, "scale": scale}


	def kp_loss(out, fut_xy, fut_valid, S, lam_ce=0.1, lam_state=0.5):
	"""fut_xy [B,H,J,2] frame units, fut_valid [B,H,J], S [B,3,2] future states (NaN censored) -> total, parts, per-mode L1 [B,M]."""
	ref, ref_ok, scale = out["ref"], out["ref_ok"], out["scale"]
	tgt = (torch.nan_to_num(fut_xy.float()) - ref[:, None]) / scale[:, None, None, None] # normalised residual target
	m = (fut_valid & ref_ok[:, None]).float()[:, :, None, :, None] # [B,H,1,J,1]
	err = ((out["res"] - tgt[:, :, None]).abs() * m).sum((1, 3, 4)) / m.sum((1, 3, 4)).clamp_min(1.0) # [B,M]
	has = (m.sum((1, 2, 3, 4)) > 0)
	best = err.detach().argmin(1)
	wta = (err.gather(1, best[:, None]).squeeze(1) * has.float()).sum() / has.float().sum().clamp_min(1.0)
	ce = F.cross_entropy(out["mode_logits"][has], best[has]) if has.any() else out["mode_logits"].sum() * 0
	sm = torch.isfinite(S)
	sbce = (F.binary_cross_entropy_with_logits(out["state_logits"], torch.nan_to_num(S), reduction="none") * sm.float()).sum() / sm.float().sum().clamp_min(1.0)
	total = wta + lam_ce * ce + lam_state * sbce
	return total, {"wta_l1": float(wta.detach()), "mode_ce": float(ce.detach()), "state_bce": float(sbce.detach())}, err


	def kpf_channel(out, anchor_ok):
	"""the distillation channel [B,KPF_DIM]: 6 state probabilities, 3 mode probabilities, valid bit (0 -> every other entry 0)."""
	p_state = torch.sigmoid(out["state_logits"]).reshape(-1, N_STATES * N_STATE_H)
	p_mode = torch.softmax(out["mode_logits"], -1)
	ok = anchor_ok.float()[:, None]
	return torch.cat([p_state * ok, p_mode * ok, ok], 1)


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