Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /tests /test_zero_init_identity.py

AsianPlayer

16 days ago

download

raw

15.2 kB

"""cwm3/tests/test_zero_init_identity.py — at initialisation the joint model's emissions equal the independent branches' BITWISE (plan §3 I).
(a) synthetic data contract (never a result): the three base branches are freshly built (seed 0); the joint model with every
    interaction (road2drv, env2veh, drv2veh), the detached K head and the pretrained direct decoders must return V, E, F and the U chain
    logits (dims 22:26) bitwise equal to the branches' own forwards (fp32, CPU); each interaction module's residual is exactly zero.
(b) if the seed-0 base checkpoints exist (BASE_DEFAULT: V0_can4s / U0_kin_can_noK / E0_pooled, seed 0), the same holds for the assembled model;
    E base pooled or region x cross-attention keys from the region bank or the encoder states are all identity at init.
(c) the naive-fusion control is NOT identity (its residual is non-zero by design) — guards against the control silently being zero-init.
(d) after one optimiser step on the interaction parameters the residuals are non-zero (the zero init is not a dead init).
Run: cd cwm3 && python -m pytest tests/test_zero_init_identity.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 pytest                                                        # 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"):
    sys.path.insert(0, str(p))
import cwm3_dataset as ds                                            # noqa: E402
import batch as bt                                                   # noqa: E402
from models import build_branch                                      # noqa: E402
from models.cwm3 import CWM3Joint, build_joint, joint_loss, BASE_DEFAULT   # noqa: E402
from models.interaction import n_params                              # noqa: E402

KEYS = sorted(set(bt.KEYS["V"]) | set(bt.KEYS["U"]) | set(bt.KEYS["E"]))
CFG = {"V": {"can_hist": "4s"}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}}


def _batch(n=8, seed=3, stats=None, e_tokens="pooled", u_cfg=None):
    """u_cfg (2026-09-11): the U batch cfg (e.g. {"cabin": "none", "kclip": 8.0} for the clip8 survivor recipe); default CFG["U"] (legacy clamp)."""
    raw = ds.synthetic_raw(64, seed=seed, mode=0, need=ds.ALL_STREAMS)
    raw["mode"][::2] = 1
    bank = ds.synthetic_bank()
    st = stats or 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, KEYS, "cpu", bank)
    idx = torch.arange(n)
    bV, tV = bt.prep_V(store, idx, st, CFG["V"])
    bU, tU = bt.prep_U(store, idx, st, u_cfg or CFG["U"])
    bE, tE = bt.prep_E(store, idx, st, {"tokens": e_tokens, "n_slots": 16})
    bR, _ = bt.prep_E(store, idx, st, CFG["R"], False)
    return {"V": bV, "U": bU, "E": bE, "R": {k: bR[k] for k in ("tok", "tok_ok", "tok_dt")}}, {"V": tV, "U": tU, "E": tE}


def _fresh_branches(seed=0, e_tokens="pooled"):
    V = build_branch("V", seed, d=256, hidden=256, dropout=0.1)
    U = build_branch("U", seed, hid=128, hidden=128, dropout=0.0, cabin="none")
    E = build_branch("E", seed, tokens=e_tokens, n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1)
    return V, U, E


def _assert_identity(model, b):
    model.eval()
    with torch.no_grad():
        o = model(b, return_trunks=True)
        oV, oU, oE = model.V(b["V"])["V"], model.U(b["U"])["U"], model.E(b["E"])
    assert torch.equal(o["V"], oV), "V emissions differ at init"
    assert torch.equal(o["U"][..., 22:], oU[..., 22:]), "U chain logits differ at init"
    assert torch.equal(o["E"], oE["E"]), "E emissions differ at init"
    assert torch.equal(o["F"], oE["F"]), "F emissions differ at init"
    assert torch.isfinite(o["U"][..., :22]).all()
    if getattr(model.U, "gaze_head", None) is not None:                # gaze U bases (2026-09-12): the gaze_off chain logits are the base's too
        with torch.no_grad():
            oG = model.U(b["U"])["G"]
        assert "G" in o and torch.equal(o["G"], oG), "gaze_off chain logits differ at init"
    return o


def test_zero_init_identity_synthetic():
    V, U, E = _fresh_branches()
    b, _ = _batch()
    m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=0)
    o = _assert_identity(m, b)
    # each interaction residual is exactly zero
    tr = o["trunks"]
    with torch.no_grad():
        assert torch.equal(m.road2drv(tr["h_u"], tr["h_tok"]), torch.zeros_like(tr["h_u"]))
        assert torch.equal(m.env2veh(torch.cat([tr["z_e"], tr["z_vf"]], 1)), torch.zeros_like(tr["z_v"]))
        assert torch.equal(m.drv2veh(tr["h_u"], b["V"]["mode"]), torch.zeros_like(tr["z_v"]))
    assert 0.7e6 <= n_params(m.road2drv) + n_params(m.env2veh) + n_params(m.drv2veh) + n_params(m.tok_embed) <= 0.9e6


@pytest.mark.parametrize("inter", [("road2drv",), ("env2veh",), ("drv2veh",), ()])
def test_single_interaction_identity(inter):
    V, U, E = _fresh_branches()
    b, _ = _batch(seed=5)
    _assert_identity(CWM3Joint(V, U, E, interactions=inter, decoder="direct", k_head=True, seed=1), b)


@pytest.mark.parametrize("e_tokens,tok_source", [("pooled", "region"), ("pooled", "enc"), ("region", "region"), ("region", "enc")])
def test_identity_e_base_and_token_source(e_tokens, tok_source):
    """E base pooled or region; cross-attention keys from the region bank (RegionTokenEmbed) or the E encoder states — identity either way."""
    V, U, E = _fresh_branches(e_tokens=e_tokens)
    b, _ = _batch(seed=7, e_tokens=e_tokens)
    m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=2, tok_source=tok_source)
    _assert_identity(m, b)
    assert (m.tok_embed is not None) == (tok_source == "region")


def test_naive_fusion_is_not_identity():
    V, U, E = _fresh_branches()
    b, _ = _batch()
    m = CWM3Joint(V, U, E, naive=True, decoder="direct", k_head=False, seed=0)
    m.eval()
    with torch.no_grad():
        o = m(b)
        oV = m.V(b["V"])["V"]
    assert not torch.equal(o["V"], oV)


def test_zero_init_is_trainable():
    V, U, E = _fresh_branches()
    b, tg = _batch()
    m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=0)
    tg["V"]["M"] = tg["V"]["M"] & torch.ones_like(tg["V"]["M"])
    opt = torch.optim.SGD(m.new_parameters(), lr=1e-2)
    m.train()
    # the detached K head sends no gradient into the GRU (only the chain losses reach it)
    out = m(b)
    _, parts = joint_loss(out, tg, {"hz": 0.5, "K": 1.0, "F": 0.1})
    parts["nll_K"].backward()
    assert all(p.grad is None or p.grad.abs().sum() == 0 for p in m.U.gru.parameters()), "K loss reached the driver GRU (head not detached)"
    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in m.k_head.parameters())
    m.zero_grad()
    out = m(b)
    loss, parts = joint_loss(out, tg, {"hz": 0.5, "K": 0.5, "F": 0.1})
    assert torch.isfinite(loss)
    loss.backward()
    assert m.road2drv.mlp.fc2.weight.grad is not None and m.road2drv.mlp.fc2.weight.grad.abs().sum() > 0
    opt.step()
    m.eval()
    with torch.no_grad():
        o2 = m(b)
        oV = m.V(b["V"])["V"]
    assert not torch.equal(o2["V"], oV)


have_base = all((CWM3 / "results" / "runs" / f"{s}_{v}_s0" / "BEST").exists() for s, v in BASE_DEFAULT.values())


@pytest.mark.skipif(not have_base, reason="seed-0 base checkpoints not on disk")
def test_zero_init_identity_base_checkpoints():
    m, info = build_joint(0, decoder="direct", k_head=True)
    st = {k: np.asarray(v) for k, v in info["stats"].items()}
    b, _ = _batch(stats=st)
    _assert_identity(m, b)
    assert 0.7e6 <= info["params"]["interaction_total"] <= 0.9e6


# ----------------------------------------------------------------------------- 2026-09-11 final assembly (roadmap §三 1, 2, 4): the survivor bases
RUNS = CWM3 / "results" / "runs"
# (stage, variant, seed-0 run dir) of the frozen configuration; every test below skips when a dir is missing (the V stats2 base is trained by the final-assembly queue)
FINAL_BASES = {"V": ("V0", "can4s_stats2", "V0_can4s_s0_stats2"), "U": ("U0", "kin_can_noK_clip8_ln", "U0_kin_can_noK_s0_kcv2_clip8_ln"), "E": ("E0", "pooled_P", "E0_pooled_P_s0_v2")}
# 2026-09-12 gaze-chain APPENDIX variant (coordinator assignment): the same frozen V / E bases with the gaze U base (recipe u_gaze_chain + u_face); runs *_final_gzfc
GZFC_BASES = dict(FINAL_BASES, U=("U0", "kin_can_noK_clip8_ln_gzfc", "U0_kin_can_noK_s0_kcv2_clip8_ln_gzfc"))
# 2026-09-13 roadmap E4 sensitivity row: the survivor U recipe on the cwm2 route-median (NON-causal) kinematics; runs *_final_ncroi
NCROI_BASES = dict(FINAL_BASES, U=("U0", "kin_can_noK_clip8_ln", "U0_kin_can_noK_s0_ncroi_clip8_ln"))
BASES_BY = {"final": FINAL_BASES, "final_gzfc": GZFC_BASES, "final_ncroi": NCROI_BASES}
_ok = lambda d: (RUNS / d / "BEST").exists() and (RUNS / d / "TRAIN_DONE").exists()          # noqa: E731


def _u_cfg(info):
    """the U batch cfg of an assembled joint (train_cwm3_joint.prep_cfgs rule): kclip + the opt-in recipe switches u_face / u_gaze_chain of the U base."""
    bc = info["base_cfgs"]["U"]
    u = {"cabin": bc.get("cabin", "none"), "kclip": bt.k_clip(bc)}
    rec = {k: bool(bc[k]) for k in ("u_face", "u_gaze_chain") if k in bc}
    if rec:
        u["recipe"] = rec
    return u


def test_branch_stats_helper():
    from models.cwm3 import branch_stats, stats_per_branch
    st = stats_per_branch([{"xmu": [1.0]}, {"xmu": [2.0], "kmu": [3.0]}, {"ymu": [4.0]}])
    st["p0_base"] = np.array([9.0], np.float32)
    u = branch_stats(st, "U")
    assert set(u) == {"xmu", "kmu", "p0_base"} and float(u["xmu"][0]) == 2.0 and float(u["p0_base"][0]) == 9.0
    plain = {"xmu": np.zeros(5)}
    assert branch_stats(plain, "V") is plain                                  # merged mode: untouched


@pytest.mark.parametrize("branch", ["U", "E", "U_gzfc"])
def test_model_kwargs_forward_survivor_configs(branch):
    """_model_kwargs forwards in_ln (U, recipe u_in_ln) / p_head (E) and base_cfg_summary carries kclip -> batch.k_clip returns 8 (not the legacy 5).
    U_gzfc (2026-09-12): the gaze U base -> gaze_chain / n_face forwarded (GRU input 18 + 9 + 1), summary carries u_gaze_chain / u_face; the plain U base gets neither."""
    import json
    from models.cwm3 import _model_kwargs, base_cfg_summary, load_base_branch
    d = (GZFC_BASES if branch == "U_gzfc" else FINAL_BASES)[branch[0]][2]
    if not _ok(d):
        pytest.skip(f"{d} not on disk")
    c = json.loads((RUNS / d / "config.json").read_text())["cfg"]
    kw, summ = _model_kwargs(c), base_cfg_summary(c)
    m, _, _, _ = load_base_branch(RUNS / d)                                   # loads the BEST state dict: fails if a head / LN is missing from the module
    if branch == "U_gzfc":
        assert kw.get("in_ln") is True and kw.get("gaze_chain") is True and kw.get("n_face") == ds.N_FACE == 9
        assert m.gaze_head is not None and m.gru.input_size == 18 + ds.N_FACE + 1
        assert summ["u_gaze_chain"] is True and summ["u_face"] is True and summ["kclip"] == 8.0
    elif branch == "U":
        assert kw.get("in_ln") is True and m.in_ln is not None and summ["u_in_ln"] is True
        assert "gaze_chain" not in kw and "n_face" not in kw and "u_gaze_chain" not in summ and "u_face" not in summ and m.gaze_head is None and m.gru.input_size == 18
        assert summ["kclip"] == 8.0 and bt.k_clip(summ) == 8.0 and bt.k_clip({"cabin": "none"}) == bt.K_CLIP_LEGACY
        assert summ["kin_source"] == "causal_routes" and summ["kin_version"] == "v2" and summ["stats_version"] == 2
    else:
        assert kw.get("p_head") is True and m.p_head is not None and summ["p_head"] is True and summ["select_LF"] is False


@pytest.mark.parametrize("which", ["final", "final_gzfc", "final_ncroi"])
@pytest.mark.parametrize("inter", [("env2veh", "drv2veh"), ("road2drv", "env2veh", "drv2veh"), ()])
def test_zero_init_identity_final_bases(which, inter):
    """the assembled final model (V stats2 / U clip8+LN causal v2 / E pooled_P) equals its bases bitwise at init — V, U chains (through the base's input
    LayerNorm), E, F AND the detached P head (copied from the E base) — with the merged statistics (the 1e-5 assert must hold on the final bases).
    final_gzfc (2026-09-12 appendix variant): the gaze U base — face channels in the GRU input (Kx 18 + 10 wide), the gaze_off chain "G" bitwise the base's too."""
    BASES = BASES_BY[which]
    if not all(_ok(v[2]) for v in BASES.values()):
        pytest.skip(f"{which} seed-0 bases not on disk")
    base = {b: (st, va) for b, (st, va, _) in BASES.items()}
    dirs = {b: str(RUNS / d) for b, (_, _, d) in BASES.items()}
    m, info = build_joint(0, base=base, dirs=dirs, stats_mode="merged", interactions=inter, decoder="direct", k_head=True)
    assert m.U.in_ln is not None and m.E.p_head is not None and m.p_head_source.startswith("copied")
    assert info["base_cfgs"]["U"]["kclip"] == 8.0 and bt.k_clip(info["base_cfgs"]["U"]) == 8.0 and info["base_cfgs"]["E"]["p_head"] is True
    assert int(info["stats"]["stats_version"]) == 2
    st = {k: np.asarray(v) for k, v in info["stats"].items()}
    b, tg = _batch(stats=st, u_cfg=_u_cfg(info))
    o = _assert_identity(m, b)
    with torch.no_grad():
        oE = m.E(b["E"])
    assert torch.equal(o["P"], oE["P"]), "P emissions differ at init (the joint P head must be the E base's copy)"
    if which == "final_gzfc":
        assert m.U.gaze_head is not None and "G" in o and o["G"].shape == (b["U"]["Kx"].shape[0], 25, 2)
        assert b["U"]["Kx"].shape[-1] == 18 + ds.N_FACE + 1 and "fmu" in st and "zG" in tg["U"]
        loss, parts = joint_loss(o, tg, {"hz": 0.5, "K": 0.5, "F": 0.1})
        assert "chain_gaze_off" in parts and torch.isfinite(loss)
    else:
        assert m.U.gaze_head is None and "G" not in o and b["U"]["Kx"].shape[-1] == 18


@pytest.mark.parametrize("run,which", [("I1_indep_s0_final", "final"), ("I1_indep_s0_final_gzfc", "final_gzfc"), ("I1_indep_s0_final_ncroi", "final_ncroi")])
def test_identity_numeric_final_run(run, which):
    """roadmap §三 4 numeric check on the real predictions: I1_indep_s0_final quick metrics == the bases' to 4 dp, layout h_on max|Δ| = 0 on every
    event a base emits, P onset probabilities identical to E0_pooled_P_s0_v2 (eval/verify_identity_final.py).  I1_indep_s0_final_gzfc (2026-09-12):
    the same against the gaze U base incl. the gaze_off_on / gaze_off_off quick metrics and the gaze_off_h_on / h_off layout block."""
    if not (RUNS / run / "PRED_DONE").exists():
        pytest.skip(f"{run} not on disk")
    sys.path.insert(0, str(CWM3 / "eval"))
    import verify_identity_final as VI
    res = VI.check(run, {b: d for b, (_, _, d) in BASES_BY[which].items()}, out=None)
    assert res["ok"], res["problems"]
    if which == "final_gzfc":
        assert any("gaze_off_on" in r[3] for r in res["rows"]) and any("gaze_off_h_on" in r[3] for r in res["rows"]), "gaze cells were not checked"

Xet Storage Details

Size:: 15.2 kB
Xet hash:: 14a690af518deb31d70b80df2a28f129e6663ea025eec36918f5eb2e3094ab25

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

	"""cwm3/tests/test_zero_init_identity.py — at initialisation the joint model's emissions equal the independent branches' BITWISE (plan §3 I).

	(a) synthetic data contract (never a result): the three base branches are freshly built (seed 0); the joint model with every
	interaction (road2drv, env2veh, drv2veh), the detached K head and the pretrained direct decoders must return V, E, F and the U chain
	logits (dims 22:26) bitwise equal to the branches' own forwards (fp32, CPU); each interaction module's residual is exactly zero.
	(b) if the seed-0 base checkpoints exist (BASE_DEFAULT: V0_can4s / U0_kin_can_noK / E0_pooled, seed 0), the same holds for the assembled model;
	E base pooled or region x cross-attention keys from the region bank or the encoder states are all identity at init.
	(c) the naive-fusion control is NOT identity (its residual is non-zero by design) — guards against the control silently being zero-init.
	(d) after one optimiser step on the interaction parameters the residuals are non-zero (the zero init is not a dead init).
	Run: cd cwm3 && python -m pytest tests/test_zero_init_identity.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 pytest # 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"):
	sys.path.insert(0, str(p))
	import cwm3_dataset as ds # noqa: E402
	import batch as bt # noqa: E402
	from models import build_branch # noqa: E402
	from models.cwm3 import CWM3Joint, build_joint, joint_loss, BASE_DEFAULT # noqa: E402
	from models.interaction import n_params # noqa: E402

	KEYS = sorted(set(bt.KEYS["V"]) \| set(bt.KEYS["U"]) \| set(bt.KEYS["E"]))
	CFG = {"V": {"can_hist": "4s"}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}}


	def _batch(n=8, seed=3, stats=None, e_tokens="pooled", u_cfg=None):
	"""u_cfg (2026-09-11): the U batch cfg (e.g. {"cabin": "none", "kclip": 8.0} for the clip8 survivor recipe); default CFG["U"] (legacy clamp)."""
	raw = ds.synthetic_raw(64, seed=seed, mode=0, need=ds.ALL_STREAMS)
	raw["mode"][::2] = 1
	bank = ds.synthetic_bank()
	st = stats or 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, KEYS, "cpu", bank)
	idx = torch.arange(n)
	bV, tV = bt.prep_V(store, idx, st, CFG["V"])
	bU, tU = bt.prep_U(store, idx, st, u_cfg or CFG["U"])
	bE, tE = bt.prep_E(store, idx, st, {"tokens": e_tokens, "n_slots": 16})
	bR, _ = bt.prep_E(store, idx, st, CFG["R"], False)
	return {"V": bV, "U": bU, "E": bE, "R": {k: bR[k] for k in ("tok", "tok_ok", "tok_dt")}}, {"V": tV, "U": tU, "E": tE}


	def _fresh_branches(seed=0, e_tokens="pooled"):
	V = build_branch("V", seed, d=256, hidden=256, dropout=0.1)
	U = build_branch("U", seed, hid=128, hidden=128, dropout=0.0, cabin="none")
	E = build_branch("E", seed, tokens=e_tokens, n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1)
	return V, U, E


	def _assert_identity(model, b):
	model.eval()
	with torch.no_grad():
	o = model(b, return_trunks=True)
	oV, oU, oE = model.V(b["V"])["V"], model.U(b["U"])["U"], model.E(b["E"])
	assert torch.equal(o["V"], oV), "V emissions differ at init"
	assert torch.equal(o["U"][..., 22:], oU[..., 22:]), "U chain logits differ at init"
	assert torch.equal(o["E"], oE["E"]), "E emissions differ at init"
	assert torch.equal(o["F"], oE["F"]), "F emissions differ at init"
	assert torch.isfinite(o["U"][..., :22]).all()
	if getattr(model.U, "gaze_head", None) is not None: # gaze U bases (2026-09-12): the gaze_off chain logits are the base's too
	with torch.no_grad():
	oG = model.U(b["U"])["G"]
	assert "G" in o and torch.equal(o["G"], oG), "gaze_off chain logits differ at init"
	return o


	def test_zero_init_identity_synthetic():
	V, U, E = _fresh_branches()
	b, _ = _batch()
	m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=0)
	o = _assert_identity(m, b)
	# each interaction residual is exactly zero
	tr = o["trunks"]
	with torch.no_grad():
	assert torch.equal(m.road2drv(tr["h_u"], tr["h_tok"]), torch.zeros_like(tr["h_u"]))
	assert torch.equal(m.env2veh(torch.cat([tr["z_e"], tr["z_vf"]], 1)), torch.zeros_like(tr["z_v"]))
	assert torch.equal(m.drv2veh(tr["h_u"], b["V"]["mode"]), torch.zeros_like(tr["z_v"]))
	assert 0.7e6 <= n_params(m.road2drv) + n_params(m.env2veh) + n_params(m.drv2veh) + n_params(m.tok_embed) <= 0.9e6


	@pytest.mark.parametrize("inter", [("road2drv",), ("env2veh",), ("drv2veh",), ()])
	def test_single_interaction_identity(inter):
	V, U, E = _fresh_branches()
	b, _ = _batch(seed=5)
	_assert_identity(CWM3Joint(V, U, E, interactions=inter, decoder="direct", k_head=True, seed=1), b)


	@pytest.mark.parametrize("e_tokens,tok_source", [("pooled", "region"), ("pooled", "enc"), ("region", "region"), ("region", "enc")])
	def test_identity_e_base_and_token_source(e_tokens, tok_source):
	"""E base pooled or region; cross-attention keys from the region bank (RegionTokenEmbed) or the E encoder states — identity either way."""
	V, U, E = _fresh_branches(e_tokens=e_tokens)
	b, _ = _batch(seed=7, e_tokens=e_tokens)
	m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=2, tok_source=tok_source)
	_assert_identity(m, b)
	assert (m.tok_embed is not None) == (tok_source == "region")


	def test_naive_fusion_is_not_identity():
	V, U, E = _fresh_branches()
	b, _ = _batch()
	m = CWM3Joint(V, U, E, naive=True, decoder="direct", k_head=False, seed=0)
	m.eval()
	with torch.no_grad():
	o = m(b)
	oV = m.V(b["V"])["V"]
	assert not torch.equal(o["V"], oV)


	def test_zero_init_is_trainable():
	V, U, E = _fresh_branches()
	b, tg = _batch()
	m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=0)
	tg["V"]["M"] = tg["V"]["M"] & torch.ones_like(tg["V"]["M"])
	opt = torch.optim.SGD(m.new_parameters(), lr=1e-2)
	m.train()
	# the detached K head sends no gradient into the GRU (only the chain losses reach it)
	out = m(b)
	_, parts = joint_loss(out, tg, {"hz": 0.5, "K": 1.0, "F": 0.1})
	parts["nll_K"].backward()
	assert all(p.grad is None or p.grad.abs().sum() == 0 for p in m.U.gru.parameters()), "K loss reached the driver GRU (head not detached)"
	assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in m.k_head.parameters())
	m.zero_grad()
	out = m(b)
	loss, parts = joint_loss(out, tg, {"hz": 0.5, "K": 0.5, "F": 0.1})
	assert torch.isfinite(loss)
	loss.backward()
	assert m.road2drv.mlp.fc2.weight.grad is not None and m.road2drv.mlp.fc2.weight.grad.abs().sum() > 0
	opt.step()
	m.eval()
	with torch.no_grad():
	o2 = m(b)
	oV = m.V(b["V"])["V"]
	assert not torch.equal(o2["V"], oV)


	have_base = all((CWM3 / "results" / "runs" / f"{s}_{v}_s0" / "BEST").exists() for s, v in BASE_DEFAULT.values())


	@pytest.mark.skipif(not have_base, reason="seed-0 base checkpoints not on disk")
	def test_zero_init_identity_base_checkpoints():
	m, info = build_joint(0, decoder="direct", k_head=True)
	st = {k: np.asarray(v) for k, v in info["stats"].items()}
	b, _ = _batch(stats=st)
	_assert_identity(m, b)
	assert 0.7e6 <= info["params"]["interaction_total"] <= 0.9e6


	# ----------------------------------------------------------------------------- 2026-09-11 final assembly (roadmap §三 1, 2, 4): the survivor bases
	RUNS = CWM3 / "results" / "runs"
	# (stage, variant, seed-0 run dir) of the frozen configuration; every test below skips when a dir is missing (the V stats2 base is trained by the final-assembly queue)
	FINAL_BASES = {"V": ("V0", "can4s_stats2", "V0_can4s_s0_stats2"), "U": ("U0", "kin_can_noK_clip8_ln", "U0_kin_can_noK_s0_kcv2_clip8_ln"), "E": ("E0", "pooled_P", "E0_pooled_P_s0_v2")}
	# 2026-09-12 gaze-chain APPENDIX variant (coordinator assignment): the same frozen V / E bases with the gaze U base (recipe u_gaze_chain + u_face); runs *_final_gzfc
	GZFC_BASES = dict(FINAL_BASES, U=("U0", "kin_can_noK_clip8_ln_gzfc", "U0_kin_can_noK_s0_kcv2_clip8_ln_gzfc"))
	# 2026-09-13 roadmap E4 sensitivity row: the survivor U recipe on the cwm2 route-median (NON-causal) kinematics; runs *_final_ncroi
	NCROI_BASES = dict(FINAL_BASES, U=("U0", "kin_can_noK_clip8_ln", "U0_kin_can_noK_s0_ncroi_clip8_ln"))
	BASES_BY = {"final": FINAL_BASES, "final_gzfc": GZFC_BASES, "final_ncroi": NCROI_BASES}
	_ok = lambda d: (RUNS / d / "BEST").exists() and (RUNS / d / "TRAIN_DONE").exists() # noqa: E731


	def _u_cfg(info):
	"""the U batch cfg of an assembled joint (train_cwm3_joint.prep_cfgs rule): kclip + the opt-in recipe switches u_face / u_gaze_chain of the U base."""
	bc = info["base_cfgs"]["U"]
	u = {"cabin": bc.get("cabin", "none"), "kclip": bt.k_clip(bc)}
	rec = {k: bool(bc[k]) for k in ("u_face", "u_gaze_chain") if k in bc}
	if rec:
	u["recipe"] = rec
	return u


	def test_branch_stats_helper():
	from models.cwm3 import branch_stats, stats_per_branch
	st = stats_per_branch([{"xmu": [1.0]}, {"xmu": [2.0], "kmu": [3.0]}, {"ymu": [4.0]}])
	st["p0_base"] = np.array([9.0], np.float32)
	u = branch_stats(st, "U")
	assert set(u) == {"xmu", "kmu", "p0_base"} and float(u["xmu"][0]) == 2.0 and float(u["p0_base"][0]) == 9.0
	plain = {"xmu": np.zeros(5)}
	assert branch_stats(plain, "V") is plain # merged mode: untouched


	@pytest.mark.parametrize("branch", ["U", "E", "U_gzfc"])
	def test_model_kwargs_forward_survivor_configs(branch):
	"""_model_kwargs forwards in_ln (U, recipe u_in_ln) / p_head (E) and base_cfg_summary carries kclip -> batch.k_clip returns 8 (not the legacy 5).
	U_gzfc (2026-09-12): the gaze U base -> gaze_chain / n_face forwarded (GRU input 18 + 9 + 1), summary carries u_gaze_chain / u_face; the plain U base gets neither."""
	import json
	from models.cwm3 import _model_kwargs, base_cfg_summary, load_base_branch
	d = (GZFC_BASES if branch == "U_gzfc" else FINAL_BASES)[branch[0]][2]
	if not _ok(d):
	pytest.skip(f"{d} not on disk")
	c = json.loads((RUNS / d / "config.json").read_text())["cfg"]
	kw, summ = _model_kwargs(c), base_cfg_summary(c)
	m, _, _, _ = load_base_branch(RUNS / d) # loads the BEST state dict: fails if a head / LN is missing from the module
	if branch == "U_gzfc":
	assert kw.get("in_ln") is True and kw.get("gaze_chain") is True and kw.get("n_face") == ds.N_FACE == 9
	assert m.gaze_head is not None and m.gru.input_size == 18 + ds.N_FACE + 1
	assert summ["u_gaze_chain"] is True and summ["u_face"] is True and summ["kclip"] == 8.0
	elif branch == "U":
	assert kw.get("in_ln") is True and m.in_ln is not None and summ["u_in_ln"] is True
	assert "gaze_chain" not in kw and "n_face" not in kw and "u_gaze_chain" not in summ and "u_face" not in summ and m.gaze_head is None and m.gru.input_size == 18
	assert summ["kclip"] == 8.0 and bt.k_clip(summ) == 8.0 and bt.k_clip({"cabin": "none"}) == bt.K_CLIP_LEGACY
	assert summ["kin_source"] == "causal_routes" and summ["kin_version"] == "v2" and summ["stats_version"] == 2
	else:
	assert kw.get("p_head") is True and m.p_head is not None and summ["p_head"] is True and summ["select_LF"] is False


	@pytest.mark.parametrize("which", ["final", "final_gzfc", "final_ncroi"])
	@pytest.mark.parametrize("inter", [("env2veh", "drv2veh"), ("road2drv", "env2veh", "drv2veh"), ()])
	def test_zero_init_identity_final_bases(which, inter):
	"""the assembled final model (V stats2 / U clip8+LN causal v2 / E pooled_P) equals its bases bitwise at init — V, U chains (through the base's input
	LayerNorm), E, F AND the detached P head (copied from the E base) — with the merged statistics (the 1e-5 assert must hold on the final bases).
	final_gzfc (2026-09-12 appendix variant): the gaze U base — face channels in the GRU input (Kx 18 + 10 wide), the gaze_off chain "G" bitwise the base's too."""
	BASES = BASES_BY[which]
	if not all(_ok(v[2]) for v in BASES.values()):
	pytest.skip(f"{which} seed-0 bases not on disk")
	base = {b: (st, va) for b, (st, va, _) in BASES.items()}
	dirs = {b: str(RUNS / d) for b, (_, _, d) in BASES.items()}
	m, info = build_joint(0, base=base, dirs=dirs, stats_mode="merged", interactions=inter, decoder="direct", k_head=True)
	assert m.U.in_ln is not None and m.E.p_head is not None and m.p_head_source.startswith("copied")
	assert info["base_cfgs"]["U"]["kclip"] == 8.0 and bt.k_clip(info["base_cfgs"]["U"]) == 8.0 and info["base_cfgs"]["E"]["p_head"] is True
	assert int(info["stats"]["stats_version"]) == 2
	st = {k: np.asarray(v) for k, v in info["stats"].items()}
	b, tg = _batch(stats=st, u_cfg=_u_cfg(info))
	o = _assert_identity(m, b)
	with torch.no_grad():
	oE = m.E(b["E"])
	assert torch.equal(o["P"], oE["P"]), "P emissions differ at init (the joint P head must be the E base's copy)"
	if which == "final_gzfc":
	assert m.U.gaze_head is not None and "G" in o and o["G"].shape == (b["U"]["Kx"].shape[0], 25, 2)
	assert b["U"]["Kx"].shape[-1] == 18 + ds.N_FACE + 1 and "fmu" in st and "zG" in tg["U"]
	loss, parts = joint_loss(o, tg, {"hz": 0.5, "K": 0.5, "F": 0.1})
	assert "chain_gaze_off" in parts and torch.isfinite(loss)
	else:
	assert m.U.gaze_head is None and "G" not in o and b["U"]["Kx"].shape[-1] == 18


	@pytest.mark.parametrize("run,which", [("I1_indep_s0_final", "final"), ("I1_indep_s0_final_gzfc", "final_gzfc"), ("I1_indep_s0_final_ncroi", "final_ncroi")])
	def test_identity_numeric_final_run(run, which):
	"""roadmap §三 4 numeric check on the real predictions: I1_indep_s0_final quick metrics == the bases' to 4 dp, layout h_on max\|Δ\| = 0 on every
	event a base emits, P onset probabilities identical to E0_pooled_P_s0_v2 (eval/verify_identity_final.py). I1_indep_s0_final_gzfc (2026-09-12):
	the same against the gaze U base incl. the gaze_off_on / gaze_off_off quick metrics and the gaze_off_h_on / h_off layout block."""
	if not (RUNS / run / "PRED_DONE").exists():
	pytest.skip(f"{run} not on disk")
	sys.path.insert(0, str(CWM3 / "eval"))
	import verify_identity_final as VI
	res = VI.check(run, {b: d for b, (_, _, d) in BASES_BY[which].items()}, out=None)
	assert res["ok"], res["problems"]
	if which == "final_gzfc":
	assert any("gaze_off_on" in r[3] for r in res["rows"]) and any("gaze_off_h_on" in r[3] for r in res["rows"]), "gaze cells were not checked"