Buckets:
| """cwm3/models/decoders.py — emission decoders of the joint model (plan §3, F): the direct 25-step head vs the chunked decoder. | |
| DirectDecoder25 : the emission head used by every cwm3 base branch so far — one flat MLP Linear(d_in -> hidden) -> GELU -> [Dropout] | |
| -> Linear(hidden -> 25*D) emitting all 25 steps at once from the trunk z (branches.flat_head; can wrap the PRETRAINED | |
| head of a base checkpoint so the joint model starts at the independent branches). | |
| ChunkedDecoder : 5 chunks x 1 s x 5 steps. Chunk c is decoded IN PARALLEL from [z ; s_c ; chunk_emb(c)] by a shared flat MLP; the latent | |
| advances DETERMINISTICALLY s_{c+1} = s_c + GRUCell([s_c ; proj(e_c)], s_c) where e_c are the EXPECTED emissions of | |
| chunk c computed from the model's own outputs (chain hazards = expected transition probabilities, MDN means, ΔK means) | |
| — never sampled physical quantities, never ground-truth futures (no teacher forcing), no posterior / KL. The decoder's | |
| forward receives only (z, expect_fn); no target tensor can reach it (tests/test_chunked_no_target_leak.py). | |
| s_0 = Linear(z). Parameter-matched to the direct head of the same branch within ±5 % by solving for `hidden` | |
| (match_hidden); d_s / d_e scale with the budget (128 / 64 for V, 64 / 32 for U and E). | |
| Expectation functions (expect_V / expect_U / expect_E) turn a chunk's raw emissions [B,5,D] into [B,5,n_expect] summaries in probability / | |
| standardised units so the GRUCell input is well scaled. No openpilot DMS field anywhere. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from . import heads as hd | |
| from .branches import flat_head | |
| H = hd.H | |
| N_CHUNKS, STEPS = 5, 5 | |
| assert N_CHUNKS * STEPS == H | |
| # ----------------------------------------------------------------------------- expected emissions (model outputs only) | |
| def expect_V(em): | |
| """[B,T,35] -> [B,T,11]: 4 chains x (h_on, h_off) + 3 MDN means (standardised units).""" | |
| v = hd.split_V(em) | |
| hz = torch.sigmoid(v["chains"]).reshape(*em.shape[:-1], -1) | |
| mu = torch.stack([hd.mdn_mean(v["steer"]), hd.mdn_mean(v["vEgo"]), hd.mdn_mean(v["aEgo"])], -1) | |
| return torch.cat([hz, mu], -1) | |
| def expect_U(em): | |
| """[B,T,26] -> [B,T,15]: 2 chains x (h_on, h_off) + 11 ΔK means (standardised).""" | |
| u = hd.split_U(em) | |
| hz = torch.sigmoid(u["chains"]).reshape(*em.shape[:-1], -1) | |
| return torch.cat([hz, u["dK_mu"]], -1) | |
| def expect_E(em): | |
| """[B,T,12] -> [B,T,12]: 6 chains x (h_on, h_off).""" | |
| return torch.sigmoid(em.float()) | |
| EXPECT = {"V": (expect_V, 11), "U": (expect_U, 15), "E": (expect_E, 12)} | |
| # ----------------------------------------------------------------------------- decoders | |
| class DirectDecoder25(nn.Module): | |
| kind = "direct" | |
| def __init__(self, d_in, hidden, D, dropout=0.0, net=None): | |
| super().__init__() | |
| self.D = D | |
| self.net = net if net is not None else flat_head(d_in, hidden, H * D, dropout) | |
| def forward(self, z, expect_fn=None): | |
| return self.net(z).view(-1, H, self.D) | |
| def gru_cell_params(d_in, d_h): | |
| return 3 * d_in * d_h + 3 * d_h * d_h + 6 * d_h | |
| def chunked_params(d_in, D, n_expect, hidden, d_s, d_e, d_chunk=16, n_chunks=N_CHUNKS, steps=STEPS): | |
| s0 = d_in * d_s + d_s | |
| emb = n_chunks * d_chunk | |
| e_proj = steps * n_expect * d_e + d_e | |
| cell = gru_cell_params(d_s + d_e, d_s) | |
| head = (d_in + d_s + d_chunk) * hidden + hidden + hidden * steps * D + steps * D | |
| return s0 + emb + e_proj + cell + head | |
| def match_hidden(d_in, D, n_expect, target, d_s, d_e, d_chunk=16, n_chunks=N_CHUNKS, steps=STEPS): | |
| """hidden width whose chunked decoder has ≈ `target` parameters (the direct head's count).""" | |
| fixed = chunked_params(d_in, D, n_expect, 0, d_s, d_e, d_chunk, n_chunks, steps) - steps * D | |
| per_h = (d_in + d_s + d_chunk) + 1 + steps * D | |
| return max(int(round((target - fixed - steps * D) / per_h)), 8) | |
| class ChunkedDecoder(nn.Module): | |
| kind = "chunked" | |
| def __init__(self, d_in, D, n_expect, hidden, d_s=64, d_e=32, d_chunk=16, dropout=0.0, n_chunks=N_CHUNKS, steps=STEPS): | |
| super().__init__() | |
| self.D, self.n_chunks, self.steps, self.d_s = D, n_chunks, steps, d_s | |
| self.s0 = nn.Linear(d_in, d_s) | |
| self.chunk_emb = nn.Embedding(n_chunks, d_chunk) | |
| self.head = flat_head(d_in + d_s + d_chunk, hidden, steps * D, dropout) | |
| self.e_proj = nn.Linear(steps * n_expect, d_e) | |
| self.cell = nn.GRUCell(d_s + d_e, d_s) | |
| def forward(self, z, expect_fn): | |
| B = z.shape[0] | |
| s = self.s0(z) | |
| outs = [] | |
| for c in range(self.n_chunks): | |
| ce = self.chunk_emb.weight[c][None].expand(B, -1) | |
| em = self.head(torch.cat([z, s, ce], 1)).view(B, self.steps, self.D) | |
| outs.append(em) | |
| if c + 1 < self.n_chunks: | |
| e = expect_fn(em).reshape(B, -1) # expected emissions of THIS chunk (model outputs) | |
| s = s + self.cell(torch.cat([s, F.gelu(self.e_proj(e))], 1), s) | |
| return torch.cat(outs, 1) # [B,25,D] | |
| def build_decoder(kind, branch, d_in, hidden_direct, dropout, pretrained_net=None, tol=0.05): | |
| """kind in {direct, direct_fresh, chunked, rssm}. -> (decoder, info dict with parameter counts). rssm (2026-09-12): models/decoders_rssm.py.""" | |
| D = {"V": hd.D_V, "U": hd.D_U, "E": hd.D_E}[branch] | |
| direct_params = (d_in * hidden_direct + hidden_direct) + (hidden_direct * H * D + H * D) | |
| if kind == "direct": | |
| dec = DirectDecoder25(d_in, hidden_direct, D, dropout, net=pretrained_net) | |
| elif kind == "direct_fresh": | |
| dec = DirectDecoder25(d_in, hidden_direct, D, dropout) | |
| elif kind == "chunked": | |
| _, n_expect = EXPECT[branch] | |
| d_s, d_e = (128, 64) if direct_params > 200_000 else (64, 32) | |
| hidden = match_hidden(d_in, D, n_expect, direct_params, d_s, d_e) | |
| dec = ChunkedDecoder(d_in, D, n_expect, hidden, d_s, d_e, dropout=dropout) | |
| n = sum(p.numel() for p in dec.parameters()) | |
| assert abs(n / direct_params - 1) <= tol, f"{branch}: chunked {n:,} vs direct {direct_params:,} ({n / direct_params - 1:+.1%})" | |
| elif kind == "rssm": # 2026-09-12 (GPU-baselines E): stochastic latent rollout, models/decoders_rssm.py | |
| from .decoders_rssm import build_rssm | |
| dec, rssm_info = build_rssm(branch, d_in, D, direct_params, dropout, tol) | |
| else: | |
| raise ValueError(kind) | |
| n = sum(p.numel() for p in dec.parameters()) | |
| hidden = dec.head[0].out_features if kind in ("chunked", "rssm") else hidden_direct | |
| return dec, {"branch": branch, "kind": kind, "params": n, "direct_params": direct_params, "ratio": n / direct_params, "hidden": hidden, | |
| **({"d_s": dec.d_s, "d_e": dec.e_proj.out_features} if kind == "chunked" else {}), **(rssm_info if kind == "rssm" else {})} | |
Xet Storage Details
- Size:
- 7.12 kB
- Xet hash:
- db66566f8d25f1fdd6f5c21e19ee6445fe8952bb7a68321685672c2b5406b3ab
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.