Buckets:
| """cwm3/tests/test_rssm_decoder.py — the RSSM stochastic rollout decoder (models/decoders_rssm.py, GPU-baselines E). | |
| (a) parameter matching: rssm vs direct within ±5 % for the three branch dims (build_decoder asserts; re-checked); | |
| (b) signature: RSSMDecoder.forward(z, expect_fn) — no target argument; the joint model's forward is unchanged; | |
| (c) eval mode never reads targets: with CONTEXT["tg"] = random targets vs None the eval output is bitwise identical under the same RNG seed; | |
| chain columns are Rao-Blackwellised probabilities in (0, 1) after sigmoid; the output has the direct decoder's shape; | |
| (d) training mode: with targets the KL term exists and is >= n_chunks x free bits (1 nat per chunk); without targets it is None and the | |
| posterior is never used; the teacher-forcing probability anneals from 1 to 0 with tf_step; | |
| (e) the joint model assembles with decoder="rssm" and its param report carries the rssm decoder info; the default decoder path is untouched | |
| (decoder="direct" still equals the branches' own emissions — the identity test's guarantee). | |
| Run: cd cwm3 && python -m pytest tests/test_rssm_decoder.py -q | |
| """ | |
| from __future__ import annotations | |
| import inspect | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train"): | |
| sys.path.insert(0, str(p)) | |
| from models import heads as hd # noqa: E402 | |
| from models.decoders import EXPECT, build_decoder # noqa: E402 | |
| from models import decoders_rssm as R # noqa: E402 | |
| from models.cwm3 import CWM3Joint # noqa: E402 | |
| from tests.test_zero_init_identity import _batch, _fresh_branches # noqa: E402 | |
| DIMS = {"V": (256 + 16, 256), "U": (128 + 16, 128), "E": (128 + 128 + 16, 256)} | |
| def test_param_match_and_signature(): | |
| for br, (d_in, hidden) in DIMS.items(): | |
| dec, info = build_decoder("rssm", br, d_in, hidden, 0.0) | |
| assert abs(info["ratio"] - 1) <= 0.05, (br, info) | |
| assert info["kind"] == "rssm" and info["n_samples"] == 32 and info["free_bits"] == 1.0 | |
| assert list(inspect.signature(R.RSSMDecoder.forward).parameters) == ["self", "z", "expect_fn"] | |
| assert list(inspect.signature(CWM3Joint.forward).parameters) == ["self", "b", "return_trunks"] | |
| def _fake_targets(B): | |
| return {"V": {"Ystd": torch.randn(B, 25, 3), "zV": (torch.rand(B, 25, 4) > 0.5).float()}, | |
| "U": {"Kstd": torch.randn(B, 25, 11), "zU": (torch.rand(B, 25, 2) > 0.5).float()}, | |
| "E": {"zD": (torch.rand(B, 25, 6) > 0.5).float()}} | |
| def test_eval_ignores_targets_and_rao_blackwellises(): | |
| for br, (d_in, hidden) in DIMS.items(): | |
| torch.manual_seed(0) | |
| dec, _ = build_decoder("rssm", br, d_in, hidden, 0.0) | |
| dec.eval() | |
| z = torch.randn(4, d_in) | |
| R.CONTEXT["tg"] = None | |
| torch.manual_seed(7) | |
| a = dec(z, EXPECT[br][0]) | |
| R.CONTEXT["tg"] = _fake_targets(4) | |
| torch.manual_seed(7) | |
| b = dec(z, EXPECT[br][0]) | |
| R.CONTEXT["tg"] = None | |
| assert torch.equal(a, b), f"{br}: eval output depends on CONTEXT targets" | |
| D = {"V": hd.D_V, "U": hd.D_U, "E": hd.D_E}[br] | |
| assert a.shape == (4, 25, D) | |
| lo, hi = R.CHAIN_COLS[br] | |
| p = torch.sigmoid(a[..., lo:hi]) | |
| assert (p > 0).all() and (p < 1).all() | |
| # Rao-Blackwellisation: the helper returns logit(mean prob) | |
| ems = torch.randn(32, 4, 25, D) | |
| rb = R.rao_blackwellise(ems, br) | |
| assert torch.allclose(torch.sigmoid(rb[..., lo:hi]), torch.sigmoid(ems[..., lo:hi]).mean(0), atol=1e-5) | |
| other = torch.ones(D, dtype=torch.bool) | |
| other[lo:hi] = False | |
| assert torch.allclose(rb[..., other], ems[..., other].mean(0), atol=1e-6) | |
| def test_training_kl_and_teacher_forcing(): | |
| br, (d_in, hidden) = "E", DIMS["E"] | |
| torch.manual_seed(1) | |
| dec, _ = build_decoder("rssm", br, d_in, hidden, 0.0) | |
| dec.train() | |
| z = torch.randn(6, d_in) | |
| R.CONTEXT["tg"] = _fake_targets(6) | |
| assert dec.p_teacher_forcing() == 1.0 | |
| em = dec(z, EXPECT[br][0]) | |
| assert em.shape == (6, 25, hd.D_E) and dec.last_kl is not None | |
| assert float(dec.last_kl) >= R.N_CHUNKS * R.FREE_BITS - 1e-4 # free bits: at least 1 nat per chunk | |
| assert dec.last_used_posterior is True # p_tf = 1 at step 0 | |
| kl = R.collect_kl() | |
| assert kl is not None and dec.last_kl is None | |
| R.CONTEXT["tg"] = None | |
| dec(z, EXPECT[br][0]) | |
| assert dec.last_kl is None and dec.last_used_posterior is False | |
| dec.tf_step.fill_(R.TF_ANNEAL_STEPS) | |
| assert dec.p_teacher_forcing() == 0.0 | |
| R.CONTEXT["tg"] = _fake_targets(6) | |
| dec(z, EXPECT[br][0]) | |
| assert dec.last_used_posterior is False and dec.last_kl is not None # free rollout, KL still regularises q -> p | |
| R.CONTEXT["tg"] = None | |
| def test_joint_assembles_with_rssm_and_direct_unchanged(): | |
| V, U, E = _fresh_branches(0) | |
| b, tg = _batch(n=4) | |
| m = CWM3Joint(V, U, E, interactions=("env2veh", "drv2veh"), decoder="rssm", seed=0).eval() | |
| rep = m.param_report() | |
| assert all(rep["decoders"][k]["kind"] == "rssm" and abs(rep["decoders"][k]["ratio"] - 1) <= 0.05 for k in ("V", "U", "E")) | |
| with torch.no_grad(): | |
| out = m(b) | |
| assert out["V"].shape == (4, 25, hd.D_V) and out["U"].shape == (4, 25, hd.D_U) and out["E"].shape == (4, 25, hd.D_E) | |
| assert torch.isfinite(out["V"]).all() and torch.isfinite(out["U"]).all() and torch.isfinite(out["E"]).all() | |
| V2, U2, E2 = _fresh_branches(0) | |
| m2 = CWM3Joint(V2, U2, E2, interactions=(), decoder="direct", seed=0).eval() | |
| with torch.no_grad(): | |
| o2 = m2(b) | |
| assert torch.equal(o2["V"], V2(b["V"])["V"]) and torch.equal(o2["E"], E2(b["E"])["E"]) | |
Xet Storage Details
- Size:
- 5.97 kB
- Xet hash:
- 230c9c8e87133d16ff2648ee09309144d9bc80b340924723b395a1e9649ebee6
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.