Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /models /branches_alt.py

AsianPlayer

16 days ago

download

raw

9.31 kB

"""cwm3/models/branches_alt.py — architecture-swap variants of the base trunks (GPU-baselines experiment B; results/tables/branch_arch_ablation.md).
Same input contract, same output contract, same flat head and mode embedding as models/branches.py; ONLY the sequence encoder of the trunk is
swapped, and its width is SOLVED so that the encoder's parameter count matches the reference trunk within ±10 % (U: GRU(18 -> 128) = 56,832;
V: TCN(6 -> 256) = 1,313,536).  Selected through the stages.yaml variant key `arch` (train_cwm3.model_kwargs forwards it; branches.build_branch
dispatches here); every existing variant has no `arch` key and builds the reference modules bit for bit.
  U arch tcn : causal TCN (branches.TCN form: k 5, dilations 1,2,4,8, width d*) -> Linear(d* -> 128)              [d* solved]
  U arch tfm : 2-layer pre-norm causal TransformerEncoder (d*, 4 heads, d_ff 2 d*) over the 40 history steps + learned step embedding, last state -> Linear(d* -> 128)
  V arch gru : GRU(6 -> h*), last state -> Linear(h* -> 256)                                                        [h* solved]
  V arch tfm : 2-layer pre-norm causal TransformerEncoder (256, 4 heads, d_ff*) + step embedding, last state (256)  [d_ff* solved]
The alternative modules subclass the reference branch, so the input LayerNorm (u_in_ln), mode embedding, flat head, cabin / adversarial arms keep the
reference code path AND the reference's seeded initial values (they are constructed before the encoder is replaced).  The joint model never uses
these variants (branch-level ablation only).  No openpilot DMS field is read anywhere (inputs: CAN whitelist, YOLO-pose kinematics).
"""
from __future__ import annotations

import torch
import torch.nn as nn

from . import heads as hd
from .branches import TCN, DriverBranch, VehicleBranch
from .identity import grad_reverse

H = hd.H
TOL = 0.10
L_HIST = 40


def n_params(m):
    return sum(p.numel() for p in m.parameters())


# ----------------------------------------------------------------------------- encoders
class CausalTransformerEncoder(nn.Module):
    """Linear(cin -> d) + learned step embedding -> n_layers pre-norm TransformerEncoder layers with a causal (upper-triangular) mask -> LN(last step)."""

    def __init__(self, cin, d, n_layers=2, n_heads=4, d_ff=None, L=L_HIST, dropout=0.0, d_out=None):
        super().__init__()
        d_ff = d_ff or 2 * d
        self.inp = nn.Linear(cin, d)
        self.pos = nn.Parameter(torch.zeros(L, d))
        nn.init.normal_(self.pos, std=0.02)
        layer = nn.TransformerEncoderLayer(d, n_heads, d_ff, dropout, activation="gelu", batch_first=True, norm_first=True)
        self.enc = nn.TransformerEncoder(layer, n_layers, enable_nested_tensor=False)
        self.norm = nn.LayerNorm(d)
        self.proj = nn.Linear(d, d_out) if (d_out is not None and d_out != d) else nn.Identity()
        self.register_buffer("mask", torch.triu(torch.ones(L, L, dtype=torch.bool), 1), persistent=False)

    def forward(self, x):                                                  # [B,L,C] -> [B,d_out]
        L = x.shape[1]
        h = self.inp(x) + self.pos[None, :L]
        h = self.enc(h, mask=self.mask[:L, :L])
        return self.proj(self.norm(h[:, -1]))


class GRUEncoder(nn.Module):
    def __init__(self, cin, hid, d_out):
        super().__init__()
        self.gru = nn.GRU(cin, hid, batch_first=True)
        self.proj = nn.Linear(hid, d_out) if hid != d_out else nn.Identity()

    def forward(self, x):
        _, h = self.gru(x)
        return self.proj(h[-1])


class TCNEncoder(nn.Module):
    def __init__(self, cin, d, d_out, dropout=0.0):
        super().__init__()
        self.tcn = TCN(cin, d, dropout)
        self.proj = nn.Linear(d, d_out) if d != d_out else nn.Identity()

    def forward(self, x):
        return self.proj(self.tcn(x))


# ----------------------------------------------------------------------------- analytic parameter counts (for the width solver)
def p_gru(cin, h, d_out):
    return 3 * (cin * h + h * h + 2 * h) + (h * d_out + d_out if h != d_out else 0)


def p_tcn(cin, d, d_out):
    return (cin * d + d) + 4 * (d * d * 5 + d) + (d * d_out + d_out if d != d_out else 0)


def p_tfm(cin, d, n_layers, d_ff, L, d_out):
    layer = (3 * d * d + 3 * d) + (d * d + d) + (d * d_ff + d_ff) + (d_ff * d + d) + 4 * d
    return (cin * d + d) + L * d + n_layers * layer + 2 * d + (d * d_out + d_out if d != d_out else 0)


def solve(fn, target, lo, hi, step=1):
    """the width in [lo, hi] (multiples of step) whose analytic count is closest to target; asserts ±TOL."""
    best, best_n = None, None
    for w in range(lo, hi + 1, step):
        n = fn(w)
        if best is None or abs(n - target) < abs(best_n - target):
            best, best_n = w, n
    assert abs(best_n / target - 1) <= TOL, f"no width within ±{TOL:.0%} of {target:,} (closest {best}: {best_n:,})"
    return best, best_n


# ----------------------------------------------------------------------------- branches
class DriverBranchAlt(DriverBranch):
    """DriverBranch with the GRU trunk replaced by a parameter-matched causal TCN (`tcn`) or causal Transformer (`tfm`)."""

    def __init__(self, arch, cin=18, hid=128, d_mode=16, cabin="none", d_cab=64, hidden=128, dropout=0.0, in_ln=False, adv_devices=0):
        super().__init__(cin, hid, d_mode, cabin, d_cab, hidden, dropout, in_ln, adv_devices)
        target = n_params(self.gru)
        del self.gru
        if arch == "tcn":
            w, n_an = solve(lambda d: p_tcn(cin, d, hid), target, 8, 256)
            self.enc = TCNEncoder(cin, w, hid, 0.0)
        elif arch == "tfm":
            w, n_an = solve(lambda d: p_tfm(cin, d, 2, 2 * d, L_HIST, hid), target, 8, 256, step=4)
            self.enc = CausalTransformerEncoder(cin, w, 2, 4, 2 * w, L_HIST, 0.0, d_out=hid)
        else:
            raise ValueError(f"U arch {arch!r}")
        n_enc = n_params(self.enc)
        assert abs(n_enc / target - 1) <= TOL, (arch, n_enc, target)
        self.arch_info = {"branch": "U", "arch": arch, "width": int(w), "encoder_params": int(n_enc), "reference_encoder_params": int(target),
                          "ratio": round(n_enc / target, 4), "reference": "GRU(18 -> 128)"}

    def forward(self, b):
        kx = b["Kx"].float()
        if self.in_ln is not None:
            kx = torch.cat([self.in_ln(kx[..., :self.N_KIN_IN]), kx[..., self.N_KIN_IN:]], -1)
        h = self.enc(kx)
        parts = [h, self.mode(b["mode"].long())]
        if self.cab is not None:
            parts.append(self.cab(b["cab"].float()) * b["cab_ok"].float()[:, None])
        out = {"U": self.head(torch.cat(parts, 1)).view(-1, H, hd.D_U)}
        if self.adv is not None:
            out["ADV"] = self.adv(grad_reverse(h))
        return out


class VehicleBranchAlt(VehicleBranch):
    """VehicleBranch with the TCN trunk replaced by a parameter-matched GRU (`gru`) or causal Transformer (`tfm`)."""

    def __init__(self, arch, d=256, d_mode=16, hidden=256, dropout=0.1, cin=6):
        super().__init__(d, d_mode, hidden, dropout, cin)
        target = n_params(self.tcn)
        del self.tcn
        if arch == "gru":
            w, n_an = solve(lambda h: p_gru(cin, h, d), target, 16, 1024)
            self.enc = GRUEncoder(cin, w, d)
        elif arch == "tfm":
            w, n_an = solve(lambda ff: p_tfm(cin, d, 2, ff, L_HIST, d), target, 16, 4096, step=16)
            self.enc = CausalTransformerEncoder(cin, d, 2, 4, w, L_HIST, dropout, d_out=d)
        else:
            raise ValueError(f"V arch {arch!r}")
        n_enc = n_params(self.enc)
        assert abs(n_enc / target - 1) <= TOL, (arch, n_enc, target)
        self.arch_info = {"branch": "V", "arch": arch, "width": int(w), "encoder_params": int(n_enc), "reference_encoder_params": int(target),
                          "ratio": round(n_enc / target, 4), "reference": "TCN(6 -> 256), k 5, dilations 1,2,4,8"}

    def forward(self, b):
        z = self.enc(b["X"].float())
        z = torch.cat([z, self.mode(b["mode"].long())], 1)
        return {"V": self.head(z).view(-1, H, hd.D_V)}


def build_alt(branch, arch, **kw):
    """called by branches.build_branch AFTER torch.manual_seed(seed) when the cfg carries `arch`."""
    if branch == "U" and arch == "driverwm":
        # 2026-09-17 (owner experiment C, pre-registration docs/decisions.md 2026-09-17T07:00:22Z): the Driver-WM temporal encoder at its published width
        # (models/u_driverwm.py; NOT parameter-matched) as the U trunk; variant kin_can_noK_clip8_ln_udwm.  Every other arch value is unchanged.
        from .u_driverwm import DriverBranchDWM
        return DriverBranchDWM(**kw)
    if branch == "U":
        return DriverBranchAlt(arch, **kw)
    if branch == "V":
        return VehicleBranchAlt(arch, **kw)
    raise ValueError(f"no architecture-swap variants for branch {branch!r} (E depth is a plain `model: {{n_layers: ...}}` variant)")


if __name__ == "__main__":
    for br, archs, kw in (("U", ("tcn", "tfm"), dict(hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True)), ("V", ("gru", "tfm"), dict(d=256, hidden=256, dropout=0.1))):
        for ar in archs:
            torch.manual_seed(0)
            m = build_alt(br, ar, **kw)
            print(br, ar, m.arch_info, "total", n_params(m))

Xet Storage Details

Size:: 9.31 kB
Xet hash:: 0b0df3cbca712cee61db5d68e1334f046d94bebbd430bb8abfd5b41e533d45bf

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

	"""cwm3/models/branches_alt.py — architecture-swap variants of the base trunks (GPU-baselines experiment B; results/tables/branch_arch_ablation.md).

	Same input contract, same output contract, same flat head and mode embedding as models/branches.py; ONLY the sequence encoder of the trunk is
	swapped, and its width is SOLVED so that the encoder's parameter count matches the reference trunk within ±10 % (U: GRU(18 -> 128) = 56,832;
	V: TCN(6 -> 256) = 1,313,536). Selected through the stages.yaml variant key `arch` (train_cwm3.model_kwargs forwards it; branches.build_branch
	dispatches here); every existing variant has no `arch` key and builds the reference modules bit for bit.
	U arch tcn : causal TCN (branches.TCN form: k 5, dilations 1,2,4,8, width d) -> Linear(d -> 128) [d* solved]
	U arch tfm : 2-layer pre-norm causal TransformerEncoder (d, 4 heads, d_ff 2 d) over the 40 history steps + learned step embedding, last state -> Linear(d* -> 128)
	V arch gru : GRU(6 -> h), last state -> Linear(h -> 256) [h* solved]
	V arch tfm : 2-layer pre-norm causal TransformerEncoder (256, 4 heads, d_ff) + step embedding, last state (256) [d_ff solved]
	The alternative modules subclass the reference branch, so the input LayerNorm (u_in_ln), mode embedding, flat head, cabin / adversarial arms keep the
	reference code path AND the reference's seeded initial values (they are constructed before the encoder is replaced). The joint model never uses
	these variants (branch-level ablation only). No openpilot DMS field is read anywhere (inputs: CAN whitelist, YOLO-pose kinematics).
	"""
	from __future__ import annotations

	import torch
	import torch.nn as nn

	from . import heads as hd
	from .branches import TCN, DriverBranch, VehicleBranch
	from .identity import grad_reverse

	H = hd.H
	TOL = 0.10
	L_HIST = 40


	def n_params(m):
	return sum(p.numel() for p in m.parameters())


	# ----------------------------------------------------------------------------- encoders
	class CausalTransformerEncoder(nn.Module):
	"""Linear(cin -> d) + learned step embedding -> n_layers pre-norm TransformerEncoder layers with a causal (upper-triangular) mask -> LN(last step)."""

	def __init__(self, cin, d, n_layers=2, n_heads=4, d_ff=None, L=L_HIST, dropout=0.0, d_out=None):
	super().__init__()
	d_ff = d_ff or 2 * d
	self.inp = nn.Linear(cin, d)
	self.pos = nn.Parameter(torch.zeros(L, d))
	nn.init.normal_(self.pos, std=0.02)
	layer = nn.TransformerEncoderLayer(d, n_heads, d_ff, dropout, activation="gelu", batch_first=True, norm_first=True)
	self.enc = nn.TransformerEncoder(layer, n_layers, enable_nested_tensor=False)
	self.norm = nn.LayerNorm(d)
	self.proj = nn.Linear(d, d_out) if (d_out is not None and d_out != d) else nn.Identity()
	self.register_buffer("mask", torch.triu(torch.ones(L, L, dtype=torch.bool), 1), persistent=False)

	def forward(self, x): # [B,L,C] -> [B,d_out]
	L = x.shape[1]
	h = self.inp(x) + self.pos[None, :L]
	h = self.enc(h, mask=self.mask[:L, :L])
	return self.proj(self.norm(h[:, -1]))


	class GRUEncoder(nn.Module):
	def __init__(self, cin, hid, d_out):
	super().__init__()
	self.gru = nn.GRU(cin, hid, batch_first=True)
	self.proj = nn.Linear(hid, d_out) if hid != d_out else nn.Identity()

	def forward(self, x):
	_, h = self.gru(x)
	return self.proj(h[-1])


	class TCNEncoder(nn.Module):
	def __init__(self, cin, d, d_out, dropout=0.0):
	super().__init__()
	self.tcn = TCN(cin, d, dropout)
	self.proj = nn.Linear(d, d_out) if d != d_out else nn.Identity()

	def forward(self, x):
	return self.proj(self.tcn(x))


	# ----------------------------------------------------------------------------- analytic parameter counts (for the width solver)
	def p_gru(cin, h, d_out):
	return 3 * (cin * h + h * h + 2 * h) + (h * d_out + d_out if h != d_out else 0)


	def p_tcn(cin, d, d_out):
	return (cin * d + d) + 4 * (d * d * 5 + d) + (d * d_out + d_out if d != d_out else 0)


	def p_tfm(cin, d, n_layers, d_ff, L, d_out):
	layer = (3 * d * d + 3 * d) + (d * d + d) + (d * d_ff + d_ff) + (d_ff * d + d) + 4 * d
	return (cin * d + d) + L * d + n_layers * layer + 2 * d + (d * d_out + d_out if d != d_out else 0)


	def solve(fn, target, lo, hi, step=1):
	"""the width in [lo, hi] (multiples of step) whose analytic count is closest to target; asserts ±TOL."""
	best, best_n = None, None
	for w in range(lo, hi + 1, step):
	n = fn(w)
	if best is None or abs(n - target) < abs(best_n - target):
	best, best_n = w, n
	assert abs(best_n / target - 1) <= TOL, f"no width within ±{TOL:.0%} of {target:,} (closest {best}: {best_n:,})"
	return best, best_n


	# ----------------------------------------------------------------------------- branches
	class DriverBranchAlt(DriverBranch):
	"""DriverBranch with the GRU trunk replaced by a parameter-matched causal TCN (`tcn`) or causal Transformer (`tfm`)."""

	def __init__(self, arch, cin=18, hid=128, d_mode=16, cabin="none", d_cab=64, hidden=128, dropout=0.0, in_ln=False, adv_devices=0):
	super().__init__(cin, hid, d_mode, cabin, d_cab, hidden, dropout, in_ln, adv_devices)
	target = n_params(self.gru)
	del self.gru
	if arch == "tcn":
	w, n_an = solve(lambda d: p_tcn(cin, d, hid), target, 8, 256)
	self.enc = TCNEncoder(cin, w, hid, 0.0)
	elif arch == "tfm":
	w, n_an = solve(lambda d: p_tfm(cin, d, 2, 2 * d, L_HIST, hid), target, 8, 256, step=4)
	self.enc = CausalTransformerEncoder(cin, w, 2, 4, 2 * w, L_HIST, 0.0, d_out=hid)
	else:
	raise ValueError(f"U arch {arch!r}")
	n_enc = n_params(self.enc)
	assert abs(n_enc / target - 1) <= TOL, (arch, n_enc, target)
	self.arch_info = {"branch": "U", "arch": arch, "width": int(w), "encoder_params": int(n_enc), "reference_encoder_params": int(target),
	"ratio": round(n_enc / target, 4), "reference": "GRU(18 -> 128)"}

	def forward(self, b):
	kx = b["Kx"].float()
	if self.in_ln is not None:
	kx = torch.cat([self.in_ln(kx[..., :self.N_KIN_IN]), kx[..., self.N_KIN_IN:]], -1)
	h = self.enc(kx)
	parts = [h, self.mode(b["mode"].long())]
	if self.cab is not None:
	parts.append(self.cab(b["cab"].float()) * b["cab_ok"].float()[:, None])
	out = {"U": self.head(torch.cat(parts, 1)).view(-1, H, hd.D_U)}
	if self.adv is not None:
	out["ADV"] = self.adv(grad_reverse(h))
	return out


	class VehicleBranchAlt(VehicleBranch):
	"""VehicleBranch with the TCN trunk replaced by a parameter-matched GRU (`gru`) or causal Transformer (`tfm`)."""

	def __init__(self, arch, d=256, d_mode=16, hidden=256, dropout=0.1, cin=6):
	super().__init__(d, d_mode, hidden, dropout, cin)
	target = n_params(self.tcn)
	del self.tcn
	if arch == "gru":
	w, n_an = solve(lambda h: p_gru(cin, h, d), target, 16, 1024)
	self.enc = GRUEncoder(cin, w, d)
	elif arch == "tfm":
	w, n_an = solve(lambda ff: p_tfm(cin, d, 2, ff, L_HIST, d), target, 16, 4096, step=16)
	self.enc = CausalTransformerEncoder(cin, d, 2, 4, w, L_HIST, dropout, d_out=d)
	else:
	raise ValueError(f"V arch {arch!r}")
	n_enc = n_params(self.enc)
	assert abs(n_enc / target - 1) <= TOL, (arch, n_enc, target)
	self.arch_info = {"branch": "V", "arch": arch, "width": int(w), "encoder_params": int(n_enc), "reference_encoder_params": int(target),
	"ratio": round(n_enc / target, 4), "reference": "TCN(6 -> 256), k 5, dilations 1,2,4,8"}

	def forward(self, b):
	z = self.enc(b["X"].float())
	z = torch.cat([z, self.mode(b["mode"].long())], 1)
	return {"V": self.head(z).view(-1, H, hd.D_V)}


	def build_alt(branch, arch, **kw):
	"""called by branches.build_branch AFTER torch.manual_seed(seed) when the cfg carries `arch`."""
	if branch == "U" and arch == "driverwm":
	# 2026-09-17 (owner experiment C, pre-registration docs/decisions.md 2026-09-17T07:00:22Z): the Driver-WM temporal encoder at its published width
	# (models/u_driverwm.py; NOT parameter-matched) as the U trunk; variant kin_can_noK_clip8_ln_udwm. Every other arch value is unchanged.
	from .u_driverwm import DriverBranchDWM
	return DriverBranchDWM(**kw)
	if branch == "U":
	return DriverBranchAlt(arch, **kw)
	if branch == "V":
	return VehicleBranchAlt(arch, **kw)
	raise ValueError(f"no architecture-swap variants for branch {branch!r} (E depth is a plain `model: {{n_layers: ...}}` variant)")


	if __name__ == "__main__":
	for br, archs, kw in (("U", ("tcn", "tfm"), dict(hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True)), ("V", ("gru", "tfm"), dict(d=256, hidden=256, dropout=0.1))):
	for ar in archs:
	torch.manual_seed(0)
	m = build_alt(br, ar, **kw)
	print(br, ar, m.arch_info, "total", n_params(m))