Buckets:
| """cwm3/tests/test_future_state_head.py — F2′ future-state head (models/future_state_head.py, train/train_future_state.py; R2 WP4 batch 1, 2026-09-16). | |
| (a) occupancy recursion: occupancy() equals models/heads.forward_marginal per state and, for constant hazards, the closed-form two-state Markov chain | |
| pi_k = pi_inf + (z0 - pi_inf)(1 - h_on - h_off)^k with pi_inf = h_on / (h_on + h_off); the transition NLL of a deterministic path is consistent with it | |
| (a path that never switches from z0 = 1 has occupancy 1 under h_off = 0). | |
| (b) strict future: leave_targets() maps horizon k (>= 1) to index k-1 of the future array (a single switch at step k* -> event at k*); an unobserved | |
| step censors at the prefix length; the INPUT features of train/train_future_state.build features (extra_features / dwell_features on the 8 s HISTORY) | |
| are bitwise unchanged when the FUTURE arrays (S / S0 / EM) are poisoned, and change when the history is perturbed. | |
| (c) first-leave survival NLL: -log p_k for an event, -log sum_{j>c} p_j when censored at c, -log p_beyond for a fully observed non-leaver; the law sums to 1; | |
| expected_leave_s of a point mass at step k is k dt. | |
| (d) losses are finite on synthetic batches, the head has < 60 k parameters and the off-task rows alone feed the remaining-dwell term. | |
| (e) make_targets follows train/batch.prep_U (z_prev = anchor state then z[:-1]; masks = onset & M & finite z / z_prev) for the two kin states. | |
| Synthetic data only (never a result). Run: cd cwm3 && python -m pytest tests/test_future_state_head.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "common", CWM3 / "eval", CWM3.parent / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| from models import heads as hd # noqa: E402 | |
| from models import future_state_head as fsh # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| import train_future_state as tfs # noqa: E402 | |
| H, S = fsh.H, len(fsh.STATES) | |
| def test_occupancy_matches_forward_marginal_and_closed_form(): | |
| g = torch.Generator().manual_seed(0) | |
| chains = torch.randn(16, H, S, 2, generator=g) | |
| z0 = (torch.rand(16, S, generator=g) > 0.5).float() | |
| pi = fsh.occupancy(chains, z0) | |
| h_on, h_off = hd.hazards(chains) | |
| for s in range(S): | |
| ref = hd.forward_marginal(h_on[:, :, s], h_off[:, :, s], z0[:, s]) | |
| assert torch.allclose(pi[:, :, s], ref, atol=1e-6) | |
| # constant hazards -> closed form | |
| hon, hoff = 0.1, 0.3 | |
| c = torch.zeros(4, H, S, 2) | |
| c[..., 0] = torch.logit(torch.tensor(hon)) | |
| c[..., 1] = torch.logit(torch.tensor(hoff)) | |
| z0 = torch.tensor([[1.0] * S, [0.0] * S, [1.0] * S, [0.0] * S]) | |
| pi = fsh.occupancy(c, z0) | |
| pinf = hon / (hon + hoff) | |
| k = torch.arange(1, H + 1).float() | |
| for b in range(4): | |
| ref = pinf + (z0[b, 0] - pinf) * (1 - hon - hoff) ** k | |
| assert torch.allclose(pi[b, :, 0], ref, atol=1e-5) | |
| # deterministic staying path: h_off -> 0 keeps occupancy at 1 from z0 = 1; its transition NLL is ~0 | |
| c2 = torch.zeros(2, H, 1, 2) | |
| c2[..., 1] = -20.0 | |
| c2[..., 0] = -20.0 | |
| pi2 = fsh.occupancy(c2, torch.ones(2, 1)) | |
| assert torch.allclose(pi2, torch.ones_like(pi2), atol=H * 2 * hd.EPS_H) # hazards are clamped to EPS_H -> pi_k >= (1 - EPS_H)^k | |
| nll = hd.chain_bce(c2[:, :, 0, 0], c2[:, :, 0, 1], torch.ones(2, H), torch.ones(2, H), torch.ones(2, H)) | |
| assert float(nll.max()) < 1e-3 | |
| def test_leave_targets_strict_future_indexing(): | |
| z0 = torch.tensor([1.0, 0.0, 1.0, 1.0, 0.0]) | |
| Sf = torch.ones(5, H) | |
| valid = torch.ones(5, H, dtype=torch.bool) | |
| Sf[0, 6:] = 0.0 # z0 = 1, first differing step k* = 7 (index 6) -> event at k = 7 | |
| Sf[1] = 0.0; Sf[1, 11] = 1.0 # z0 = 0, switch at step 12 | |
| # row 2: never leaves -> k = H, no event ('beyond H') | |
| valid[3, 4:] = False # row 3: observed steps 1..4 only, no switch -> censored at c = 4 | |
| Sf[4] = 0.0; valid[4, 2] = False; Sf[4, 5] = 1.0 # row 4: gap at step 3 -> prefix = 2 steps, the later switch is NOT used (strict prefix) | |
| ev, k = fsh.leave_targets(z0, Sf, valid) | |
| assert ev.tolist() == [True, True, False, False, False] | |
| assert k.tolist() == [7, 12, H, 4, 2] | |
| # poisoning the future beyond the first switch changes nothing before it | |
| Sf2 = Sf.clone(); Sf2[0, 10:] = torch.rand(H - 10) | |
| ev2, k2 = fsh.leave_targets(z0, Sf2, valid) | |
| assert ev2[0] == ev[0] and k2[0] == k[0] | |
| def test_leave_nll_forms_and_expectation(): | |
| g = torch.Generator().manual_seed(1) | |
| logits = torch.randn(4, fsh.N_BINS, generator=g) | |
| p = fsh.leave_law(logits) | |
| assert torch.allclose(p.sum(-1), torch.ones(4), atol=1e-6) | |
| ev = torch.tensor([True, False, False, True]) | |
| k = torch.tensor([7, H, 4, 1]) | |
| m = torch.ones(4) | |
| l = fsh.leave_nll(logits, ev, k, m) | |
| assert torch.allclose(l[0], -p[0, 6].log(), atol=1e-5) # event at step 7 -> bin 6 | |
| assert torch.allclose(l[1], -p[1, H].log(), atol=1e-5) # fully observed non-leaver -> beyond bin | |
| assert torch.allclose(l[2], -p[2, 4:].sum().log(), atol=1e-5) # censored at c = 4 -> mass of steps >= 5 incl. beyond | |
| assert torch.allclose(l[3], -p[3, 0].log(), atol=1e-5) | |
| assert float(fsh.leave_nll(logits, ev, k, torch.zeros(4)).abs().max()) == 0.0 | |
| point = torch.full((1, fsh.N_BINS), -30.0); point[0, 9] = 30.0 # step 10 | |
| assert abs(float(fsh.expected_leave_s(fsh.leave_law(point))) - 10 * fsh.DT) < 1e-4 | |
| assert abs(float(fsh.leave_within(fsh.leave_law(point), 10)) - 1.0) < 1e-4 and float(fsh.leave_within(fsh.leave_law(point), 9)) < 1e-4 | |
| def _synthetic_batch(B=64, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| z = torch.randn(B, 128, generator=g) | |
| dw = torch.randn(B, len(tfs.FEATURES), generator=g) | |
| mode = (torch.rand(B, generator=g) > 0.5).long() | |
| S0 = (torch.rand(B, S, generator=g) > 0.5).float() | |
| Sf = S0[:, None, :].repeat(1, H, 1) | |
| flip = torch.rand(B, H, S, generator=g) < 0.05 | |
| Sf = torch.where(flip.cumsum(1) > 0, 1 - S0[:, None, :], Sf) | |
| Sm = torch.rand(B, H, S, generator=g) > 0.05 | |
| z_prev = torch.cat([S0[:, None], Sf[:, :-1]], 1) | |
| tg = {"z_prev": z_prev, "z": Sf, "m": Sm, "Sfut": Sf, "Sm": Sm, "z0": S0, "z0valid": torch.ones(B, S, dtype=torch.bool)} | |
| return z, dw, mode, tg | |
| def test_head_loss_finite_params_and_rem_dwell_rows(): | |
| torch.manual_seed(0) | |
| m = fsh.FutureStateHead(d_z=128, d_dwell=len(tfs.FEATURES), n_states=S) | |
| assert fsh.n_params(m) < 60_000, fsh.n_params(m) | |
| z, dw, mode, tg = _synthetic_batch() | |
| out = m(z, dw, mode) | |
| assert out["chains"].shape == (64, H, S, 2) and out["first_leave"].shape == (64, S, fsh.N_BINS) | |
| loss, parts = fsh.future_state_loss(out, tg) | |
| assert torch.isfinite(loss) and all(torch.isfinite(torch.as_tensor(v)) for v in parts.values()) | |
| assert set(k for k in parts if k.startswith("occ_") and k != "occ_bce") == {f"occ_{s}" for s in fsh.STATES} | |
| # remaining-dwell term ignores on-task rows: force every anchor on-task -> rem_* = 0 | |
| tg2 = dict(tg) | |
| tg2["z0"] = torch.tensor([[1.0 - fsh.OFF_VALUE[s] for s in fsh.STATES]]).repeat(64, 1) | |
| _, parts2 = fsh.future_state_loss(out, tg2) | |
| assert all(float(parts2[f"rem_{s}"]) == 0.0 for s in fsh.STATES) | |
| res = fsh.future_state_outputs(out, tg["z0"]) | |
| assert res["pi"].shape == (64, H, S) and torch.all(res["ED_off_s"] >= 0) and torch.all(res["ED_off_s"] <= H * fsh.DT + 1e-4) | |
| loss.backward() | |
| assert all(p.grad is not None for p in m.parameters()) | |
| def test_features_read_history_only(): | |
| """build inputs from a synthetic 8 s history; poisoning the FUTURE (targets) cannot move them; perturbing the history does.""" | |
| rng = np.random.default_rng(0) | |
| n = 32 | |
| Kh8 = rng.normal(0, 1, (n, 80, 11)).astype(np.float32) | |
| Kh8[:, :, tfs.wfe.IDX_MIN_DIST] = np.abs(Kh8[:, :, tfs.wfe.IDX_MIN_DIST]) * 0.2 | |
| V = rng.random((n, 80)) > 0.05 | |
| X1 = tfs.extra_features(Kh8, V, 0.19, 0.10) | |
| X1d = tfs.wfe.dwell_features(Kh8, V, 0.19, 0.10) | |
| assert X1.shape == (n, 4) and X1d.shape == (n, 11) | |
| # the feature builders take no future argument at all: poison "future" arrays and recompute -> identical | |
| S_future = rng.random((n, 25, 2)) # noqa: F841 (never an input) | |
| X2 = tfs.extra_features(Kh8.copy(), V.copy(), 0.19, 0.10) | |
| assert np.array_equal(np.nan_to_num(X1, nan=-9), np.nan_to_num(X2, nan=-9)) | |
| Kh8p = Kh8.copy(); Kh8p[:, -1, tfs.wfe.IDX_MIN_DIST] = 5.0 # hands leave the wheel at the anchor step | |
| X3 = tfs.extra_features(Kh8p, V, 0.19, 0.10) | |
| assert not np.array_equal(np.nan_to_num(X1, nan=-9), np.nan_to_num(X3, nan=-9)) | |
| # run length: a state constant over the whole window -> cap; a switch 1.5 s ago -> 1.5 s | |
| st = np.zeros((2, 80), np.float32); st[1, -15:] = 1.0 | |
| rl = tfs._run_length_end(st, np.ones((2, 80), bool)) | |
| assert abs(rl[0] - tfs.wfe.DWELL_CAP) < 1e-6 and abs(rl[1] - 1.5) < 1e-6 | |
| def test_make_targets_follows_prep_u(): | |
| raw = tfs.ds.synthetic_raw(48, seed=2, mode=0, need=tfs.ds.ALL_STREAMS) | |
| zu = {"S": raw["S"], "S0": raw["S0"].astype(np.float32), "EM": raw["EM"], "M": raw["M"], "Sg": raw["Sg"], "Sg0": raw["Sg0"].astype(np.float32), "EMg": raw["EMg"]} | |
| tg = tfs.make_targets(zu, "cpu") | |
| st = bt.compute_stats(raw, {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)}) | |
| store = bt.GpuStore(raw, bt.KEYS["U"], "cpu") | |
| _, ref = bt.prep_U(store, torch.arange(48), st, {"cabin": "none", "kclip": 8.0, "recipe": {"u_gaze_chain": True}}, True) | |
| assert torch.equal(tg["z"][..., :2], ref["zU"]) and torch.equal(tg["z_prev"][..., :2], ref["zU_prev"]) and torch.equal(tg["m"][..., :2], ref["zUm"]) | |
| assert torch.equal(tg["z0"][:, :2], ref["z0"]) and torch.equal(tg["z0valid"][:, :2], ref["z0valid"]) | |
| assert torch.equal(tg["z"][..., 2], ref["zG"]) and torch.equal(tg["m"][..., 2], ref["zGm"]) and torch.equal(tg["z0"][:, 2], ref["zG0"]) | |
Xet Storage Details
- Size:
- 10.5 kB
- Xet hash:
- 8d310fb2ac21cd6910a57b59e5d53988e958a13b0994991d4cb50e657258da45
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.