Buckets:
| """cwm3/models/decoders_rssm.py — RSSM-style STOCHASTIC latent rollout decoder: the cwm3-native "world-model" baseline (GPU-baselines experiment E; | |
| results/tables/stage_F_rssm.md). Drop-in decoder kind `rssm` for models/decoders.build_decoder (the joint model, models/cwm3.py, is untouched). | |
| Per branch, from the trunk z_in [B,d_in]: s_0 = tanh(W z_in); for chunk c = 0..4 (1 s each, 5 steps): | |
| prior p(z_c | s_c) = N(mu_p, sig_p) (Linear(d_s -> 2 d_z)) | |
| posterior q(z_c | s_c, y_c) = N(mu_q, sig_q) (Linear(d_s + d_y -> 2 d_z); y_c = the TRUE targets of chunk c, TRAINING ONLY) | |
| emission em_c = head([z_in ; s_c ; z_c ; chunk_emb_c]) -> [B,5,D] (flat_head, width solved so the decoder matches the direct head ±5 %) | |
| transition s_{c+1} = GRUCell([z_c ; gelu(e_proj(expect_fn(em_c)))], s_c) (the model's own expected emissions feed back, as in the chunked decoder) | |
| Training (module in .train()): the latent z_c is the posterior sample with probability p_tf (teacher forcing) and the prior sample otherwise; | |
| p_tf = max(0, 1 - step / TF_ANNEAL_STEPS) anneals from 1 to 0 over the first TF_ANNEAL_STEPS optimiser steps (counted in the `tf_step` buffer, | |
| saved with the checkpoint) -> pure free rollout afterwards. KL term = sum_c max(KL(q_c || p_c), FREE_BITS) (1 nat of free bits per chunk), | |
| batch mean, stored in `last_kl` and added to the joint loss by train/train_cwm3_rssm.py (lambda KL, default 1); it is never part of the | |
| selection objective (part name `kl_rssm`). The targets reach the decoder ONLY through CONTEXT["tg"], which the wrapper trainer fills from | |
| prep_joint while the model is in training mode; in .eval() the decoder ignores CONTEXT entirely (tests/test_rssm_decoder.py). | |
| Prediction (module in .eval()): N_SAMPLES (32) independent prior rollouts; Rao-Blackwellised emissions — chain logits are returned as | |
| logit(mean_s sigmoid(logit_s)) so the downstream sigmoid yields the sample-averaged hazard (the cwm2 rollout convention); the MDN / Gaussian | |
| parameter columns are averaged over samples (an approximation: the mixture over samples is not representable in the fixed MDN layout; | |
| the onset cells, which decide the claim, are exact). forward(z, expect_fn) has NO target argument, like ChunkedDecoder. | |
| Parameter matching: match_hidden_rssm solves the head width for the direct head's count; build_rssm asserts |ratio - 1| <= 5 %. | |
| No openpilot DMS field is read anywhere (inputs: the branch trunks; training targets: the CAN / kinematics / demand chains of the packs). | |
| """ | |
| from __future__ import annotations | |
| import math | |
| 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 | |
| N_SAMPLES, FREE_BITS, TF_ANNEAL_STEPS = 32, 1.0, 4000 | |
| CONTEXT = {"tg": None} # {"V": tgV, "U": tgU, "E": tgE} of the CURRENT training batch (set by train_cwm3_rssm.prep_joint wrapper; None otherwise) | |
| REGISTRY = [] # live RSSM decoders of this process (collect_kl) | |
| # chain-logit columns of each branch's emission layout (heads.py); everything else is a distribution parameter averaged over samples | |
| CHAIN_COLS = {"V": (27, hd.D_V), "U": (22, hd.D_U), "E": (0, hd.D_E)} | |
| # target features per step for the posterior: V = [Ystd(3) ; zV(4)], U = [Kstd(11) ; zU(2)], E = [zD(6)] | |
| D_Y_STEP = {"V": 3 + len(hd.V_CHAINS), "U": 11 + len(hd.U_CHAINS), "E": len(hd.E_CHAINS)} | |
| def target_feats(branch, tg): | |
| """the branch's training targets -> [B, N_CHUNKS, STEPS * D_Y_STEP] (chunk-major); None when no targets are in CONTEXT.""" | |
| if tg is None or branch not in tg or tg[branch] is None: | |
| return None | |
| t = tg[branch] | |
| if branch == "V": | |
| y = torch.cat([t["Ystd"].float(), t["zV"].float()], -1) | |
| elif branch == "U": | |
| y = torch.cat([t["Kstd"].float(), t["zU"].float()], -1) | |
| else: | |
| y = t["zD"].float() | |
| B = y.shape[0] | |
| return torch.nan_to_num(y).reshape(B, N_CHUNKS, STEPS * y.shape[-1]) | |
| def gaussian_kl(mq, lq, mp, lp): | |
| """KL(N(mq, e^lq) || N(mp, e^lp)) per dimension.""" | |
| return lp - lq + (torch.exp(2 * lq) + (mq - mp) ** 2) / (2 * torch.exp(2 * lp)) - 0.5 | |
| def rao_blackwellise(ems, branch, eps=hd.EPS_H): | |
| """ems [S,B,T,D] sample emissions -> [B,T,D]: chain columns = logit(mean_s sigmoid), other columns = mean_s.""" | |
| ems = ems.float() | |
| lo, hi = CHAIN_COLS[branch] | |
| out = ems.mean(0) | |
| p = torch.sigmoid(ems[..., lo:hi]).mean(0).clamp(eps, 1 - eps) | |
| out[..., lo:hi] = torch.log(p) - torch.log1p(-p) | |
| return out | |
| def rssm_params(d_in, D, n_expect, d_y, hidden, d_s, d_z, d_e, d_chunk=16, n_chunks=N_CHUNKS, steps=STEPS): | |
| s0 = d_in * d_s + d_s | |
| emb = n_chunks * d_chunk | |
| prior = d_s * 2 * d_z + 2 * d_z | |
| post = (d_s + d_y) * 2 * d_z + 2 * d_z | |
| e_proj = steps * n_expect * d_e + d_e | |
| cell = 3 * (d_z + d_e) * d_s + 3 * d_s * d_s + 6 * d_s | |
| head = (d_in + d_s + d_z + d_chunk) * hidden + hidden + hidden * steps * D + steps * D | |
| return s0 + emb + prior + post + e_proj + cell + head | |
| def match_hidden_rssm(d_in, D, n_expect, d_y, target, d_s, d_z, d_e, d_chunk=16): | |
| fixed = rssm_params(d_in, D, n_expect, d_y, 0, d_s, d_z, d_e, d_chunk) - STEPS * D | |
| per_h = (d_in + d_s + d_z + d_chunk) + 1 + STEPS * D | |
| return max(int(round((target - fixed - STEPS * D) / per_h)), 8) | |
| class RSSMDecoder(nn.Module): | |
| kind = "rssm" | |
| def __init__(self, branch, d_in, D, n_expect, d_y, hidden, d_s=64, d_z=16, d_e=32, d_chunk=16, dropout=0.0, n_chunks=N_CHUNKS, steps=STEPS, | |
| n_samples=N_SAMPLES, free_bits=FREE_BITS, tf_anneal_steps=TF_ANNEAL_STEPS): | |
| super().__init__() | |
| self.branch, self.D, self.n_chunks, self.steps, self.d_s, self.d_z, self.d_y = branch, D, n_chunks, steps, d_s, d_z, d_y | |
| self.n_samples, self.free_bits, self.tf_anneal_steps = int(n_samples), float(free_bits), int(tf_anneal_steps) | |
| self.s0 = nn.Sequential(nn.Linear(d_in, d_s), nn.Tanh()) | |
| self.chunk_emb = nn.Embedding(n_chunks, d_chunk) | |
| self.prior = nn.Linear(d_s, 2 * d_z) | |
| self.post = nn.Linear(d_s + d_y, 2 * d_z) | |
| self.e_proj = nn.Linear(steps * n_expect, d_e) | |
| self.cell = nn.GRUCell(d_z + d_e, d_s) | |
| self.head = flat_head(d_in + d_s + d_z + d_chunk, hidden, steps * D, dropout) | |
| self.register_buffer("tf_step", torch.zeros((), dtype=torch.long)) | |
| self.last_kl = None | |
| self.last_used_posterior = None | |
| REGISTRY.append(self) | |
| def p_teacher_forcing(self): | |
| return max(0.0, 1.0 - float(self.tf_step) / max(self.tf_anneal_steps, 1)) | |
| def _rollout(self, z_in, expect_fn, y=None, use_post=False): | |
| """one stochastic rollout -> (emissions [B,25,D], kl scalar or None); y [B,n_chunks,d_y] enables the posterior / KL.""" | |
| B = z_in.shape[0] | |
| s = self.s0(z_in) | |
| outs, kls = [], [] | |
| for c in range(self.n_chunks): | |
| mp, lp = self.prior(s).float().chunk(2, -1) | |
| lp = lp.clamp(-6.0, 2.0) | |
| if y is not None: | |
| mq, lq = self.post(torch.cat([s, y[:, c].to(s.dtype)], 1)).float().chunk(2, -1) | |
| lq = lq.clamp(-6.0, 2.0) | |
| kls.append(torch.clamp(gaussian_kl(mq, lq, mp, lp).sum(-1), min=self.free_bits)) # free bits per chunk | |
| eps = torch.randn(B, self.d_z, device=z_in.device, dtype=torch.float32) | |
| zc = (mq + lq.exp() * eps) if (y is not None and use_post) else (mp + lp.exp() * eps) | |
| zc = zc.to(s.dtype) | |
| ce = self.chunk_emb.weight[c][None].expand(B, -1) | |
| em = self.head(torch.cat([z_in, s, zc, ce], 1)).view(B, self.steps, self.D) | |
| outs.append(em) | |
| if c + 1 < self.n_chunks: | |
| e = expect_fn(em).reshape(B, -1).to(s.dtype) | |
| s = self.cell(torch.cat([zc, F.gelu(self.e_proj(e))], 1), s) | |
| kl = torch.stack(kls, 1).sum(1).mean() if kls else None | |
| return torch.cat(outs, 1), kl | |
| def forward(self, z, expect_fn): | |
| if self.training: | |
| y = target_feats(self.branch, CONTEXT.get("tg")) | |
| use_post = bool(y is not None and torch.rand(()) < self.p_teacher_forcing()) | |
| self.tf_step += 1 | |
| em, kl = self._rollout(z, expect_fn, y, use_post) | |
| self.last_kl, self.last_used_posterior = kl, use_post | |
| return em | |
| self.last_kl = None | |
| ems = torch.stack([self._rollout(z, expect_fn)[0] for _ in range(self.n_samples)], 0) # eval: prior only, never CONTEXT | |
| return rao_blackwellise(ems, self.branch).to(z.dtype if z.dtype != torch.float16 else torch.float32) | |
| def collect_kl(): | |
| """sum of the live decoders' last training-batch KL terms (None -> 0); resets them.""" | |
| tot = None | |
| for d in REGISTRY: | |
| if d.last_kl is not None: | |
| tot = d.last_kl if tot is None else tot + d.last_kl | |
| d.last_kl = None | |
| return tot | |
| def build_rssm(branch, d_in, D, direct_params, dropout=0.0, tol=0.05): | |
| from .decoders import EXPECT | |
| _, n_expect = EXPECT[branch] | |
| d_y = STEPS * D_Y_STEP[branch] | |
| d_s, d_z, d_e = (128, 32, 64) if direct_params > 200_000 else (64, 16, 32) | |
| hidden = match_hidden_rssm(d_in, D, n_expect, d_y, direct_params, d_s, d_z, d_e) | |
| dec = RSSMDecoder(branch, d_in, D, n_expect, d_y, hidden, d_s, d_z, d_e, dropout=dropout) | |
| n = sum(p.numel() for p in dec.parameters()) | |
| assert abs(n / direct_params - 1) <= tol, f"{branch}: rssm {n:,} vs direct {direct_params:,} ({n / direct_params - 1:+.1%})" | |
| return dec, {"hidden": hidden, "d_s": d_s, "d_z": d_z, "d_e": d_e, "d_y": d_y, "n_samples": dec.n_samples, "free_bits": dec.free_bits, "tf_anneal_steps": dec.tf_anneal_steps} | |
Xet Storage Details
- Size:
- 9.78 kB
- Xet hash:
- 05d67cb341aeaaddf206284a505390caf343717044b0bcb6bd0359970a6c3ced
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.