Buckets:
| """cwm3/train/mtl.py — multi-task balancing block of the joint trainer (WP3 reviewer control, pre-registered 2026-09-16T06:46:35Z; train_cwm3_joint --mtl). | |
| The joint model's training loss is the PLAIN SUM of the three branch losses (vehicle L_V, driver L_U, road L_E; models/cwm3.joint_loss) plus the | |
| detached auxiliary chains. Reviewers asked whether standard multi-task optimisation changes the picture (Kurin et al. 2022 / Xin et al. 2022 report | |
| that plain summation is hard to beat). Three methods, each applied to the three BRANCH losses (the detached P / R heads keep their plain terms — | |
| they share no parameter with any branch): | |
| uw : Kendall, Gal & Cipolla (2018) homoscedastic-uncertainty weighting — one learned log-variance s_i per branch, objective | |
| sum_i exp(-s_i) * (L_i - c) + s_i; the three scalars live in the MAIN optimiser (weight decay 0). | |
| famo : Liu, Liu, Stone & Liu (2023) Fast Adaptive Multitask Optimization, Algorithm 1 — logits xi, z = softmax(xi); the model step follows | |
| sum_i (z_i / D_i) grad L_i (normalised by sum_i z_i / D_i, D_i = L_i - c); xi is updated from the log-loss improvement on the SAME batch | |
| after the model step (one extra no-grad forward), with its own Adam (beta = 0.025, gamma = 0.001 = weight decay), as in the paper / official code. | |
| pcgrad : Yu et al. (2020) gradient surgery — per-branch gradients on the SHARED parameters; for every task i and every other task j in random order, | |
| if g_i . g_j < 0 then g_i <- g_i - (g_i . g_j / |g_j|^2) g_j; the shared gradient is the sum of the projected gradients, every task-specific | |
| parameter keeps its own task's gradient (= the plain-sum gradient). Shared = parameters connected in the autograd graph to >= 2 branch losses, | |
| determined in a PROBE step (torch.autograd.grad with allow_unused: a parameter that receives a gradient tensor — even an exactly-zero one — from two | |
| losses is shared; the zero-initialised interaction modules make the E / U trunks' vehicle-loss gradients EXACTLY zero at step 0, so a magnitude | |
| criterion would miss them while they are structurally shared; the probe also records how many shared parameters carried non-zero gradients). | |
| Loss shift c (documented adaptation): the vehicle branch contains continuous-density NLL terms (MDN steer / vEgo / aEgo, Gaussian K) that are NEGATIVE | |
| (≈ -3.8 on the final bases), so the likelihood-based weightings (uw needs L_i >= 0 for a finite optimum; famo takes log L_i) are applied to | |
| D_i = L_i - c with the FIXED constant c = -10, a bound below every batch loss observed in the final runs (train totals -1.4 ... +4.3). The shift is a | |
| constant: it changes no model gradient, only the weight dynamics. Gradient clipping and the optimiser are those of the recipe, unchanged. | |
| `kind == "none"` is the identity: the trainer's original path (total.backward()) runs untouched — tests/test_mtl_joint.py asserts bit-identical losses. | |
| CPU / GPU agnostic; no data access; no openpilot DMS field. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| KINDS = ("none", "uw", "famo", "pcgrad") | |
| TASKS = ("V", "U", "E") | |
| LOSS_SHIFT = -10.0 | |
| class MTL: | |
| def __init__(self, kind="none", device="cpu", seed=0, shift=LOSS_SHIFT, famo_lr=0.025, famo_gamma=0.001): | |
| assert kind in KINDS, kind | |
| self.kind, self.shift, self.n = kind, float(shift), len(TASKS) | |
| self.device = torch.device(device) | |
| self.gen = torch.Generator().manual_seed(int(seed) + 4242) # PCGrad task order (CPU generator: deterministic on every device) | |
| self.log_var = None # uw: 3 learned log-variances (main optimiser) | |
| self.xi, self.xi_opt, self._famo_prev = None, None, None # famo: logits, their Adam, the pre-step losses of the current batch | |
| self.shared, self.probe_info = None, {} # pcgrad: shared parameters (probe) + report | |
| self.last_weights = [1.0] * self.n | |
| if kind == "uw": | |
| self.log_var = torch.nn.Parameter(torch.zeros(self.n, device=self.device)) | |
| if kind == "famo": | |
| self.xi = torch.zeros(self.n, device=self.device, requires_grad=True) | |
| self.xi_opt = torch.optim.Adam([self.xi], lr=float(famo_lr), weight_decay=float(famo_gamma)) | |
| # ------------------------------------------------------------------ optimiser / checkpoint plumbing | |
| def optim_params(self): | |
| """parameters the MAIN optimiser must own (uw only; weight decay 0 for them).""" | |
| return [self.log_var] if self.kind == "uw" else [] | |
| def state_dict(self): | |
| sd = {"kind": self.kind, "shift": self.shift, "last_weights": list(self.last_weights), "probe_info": dict(self.probe_info)} | |
| if self.kind == "uw": | |
| sd["log_var"] = self.log_var.detach().cpu().clone() | |
| if self.kind == "famo": | |
| sd["xi"] = self.xi.detach().cpu().clone() | |
| sd["xi_opt"] = self.xi_opt.state_dict() | |
| return sd | |
| def load_state_dict(self, sd): | |
| assert sd.get("kind", "none") == self.kind, (sd.get("kind"), self.kind) | |
| self.last_weights = list(sd.get("last_weights", self.last_weights)) | |
| self.probe_info = dict(sd.get("probe_info", {})) | |
| if self.kind == "uw": | |
| with torch.no_grad(): | |
| self.log_var.copy_(sd["log_var"].to(self.device)) | |
| if self.kind == "famo": | |
| with torch.no_grad(): | |
| self.xi.copy_(sd["xi"].to(self.device)) | |
| self.xi_opt.load_state_dict(sd["xi_opt"]) | |
| # ------------------------------------------------------------------ weights | |
| def weights(self): | |
| """the effective per-branch weights of the last step (report only).""" | |
| return [float(w) for w in self.last_weights] | |
| def _shifted(self, branches): | |
| return [(branches[t] - self.shift).clamp_min(1e-3) for t in TASKS] | |
| # ------------------------------------------------------------------ the step | |
| def probe_shared(self, branches, params): | |
| """PCGrad: which of `params` (trainable) are connected to >= 2 branch losses. Called once per training chunk on the first batch (retain_graph).""" | |
| conn = torch.zeros(len(params), dtype=torch.int32) | |
| nonzero = torch.zeros(len(params), dtype=torch.int32) | |
| for t in TASKS: | |
| g = torch.autograd.grad(branches[t], params, retain_graph=True, allow_unused=True) | |
| for i, gi in enumerate(g): | |
| if gi is not None: | |
| conn[i] += 1 | |
| if bool((gi != 0).any()): | |
| nonzero[i] += 1 | |
| self.shared = [p for p, c in zip(params, conn.tolist()) if c >= 2] | |
| self.probe_info = {"n_params_trainable": int(len(params)), "n_shared_tensors": int(len(self.shared)), | |
| "n_shared_numel": int(sum(p.numel() for p in self.shared)), | |
| "n_shared_nonzero_grad_tensors": int(sum(1 for c, z in zip(conn.tolist(), nonzero.tolist()) if c >= 2 and z >= 2)), | |
| "connectivity_hist": {str(k): int((conn == k).sum()) for k in range(self.n + 1)}} | |
| return self.shared | |
| def backward(self, total, branches, params=None): | |
| """Backward pass of one batch. total = the plain joint loss (graph attached), branches = {"V","U","E","aux"} scalar tensors of the same graph. | |
| none : total.backward() (the trainer's original path — kept in the trainer itself; this method mirrors it for tests) | |
| uw : (sum_i exp(-s_i) D_i + s_i + aux).backward() | |
| famo : (sum_i z_i log D_i / c + aux).backward(), c = sum_i z_i / D_i (detached); remembers D_i for the post-step update | |
| pcgrad: per-task gradients on self.shared (probe_shared must have run), total.backward() for every parameter, then the projected sum | |
| overwrites .grad of the shared parameters. | |
| -> the scalar that was back-propagated (for logging).""" | |
| if self.kind == "none": | |
| total.backward() | |
| return total.detach() | |
| if self.kind == "uw": | |
| D = self._shifted(branches) | |
| w = torch.exp(-self.log_var) | |
| loss = sum(w[i] * D[i] + self.log_var[i] for i in range(self.n)) + branches["aux"] | |
| self.last_weights = w.detach().cpu().tolist() | |
| loss.backward() | |
| return loss.detach() | |
| if self.kind == "famo": | |
| D = self._shifted(branches) | |
| z = torch.softmax(self.xi, -1) | |
| Dd = torch.stack([d.detach() for d in D]) | |
| c = (z.detach() / Dd).sum() | |
| loss = sum((z[i].detach() / c) * torch.log(D[i]) for i in range(self.n)) + branches["aux"] | |
| self._famo_prev = Dd.clone() | |
| self.last_weights = (z.detach() / (c * Dd)).cpu().tolist() # effective coefficient of grad L_i | |
| loss.backward() | |
| return loss.detach() | |
| # pcgrad | |
| assert self.shared is not None, "probe_shared() must run before the first PCGrad step" | |
| shared = self.shared | |
| if not shared: # nothing shared (e.g. the `none` interaction set): the plain sum IS the PCGrad update | |
| total.backward() | |
| return total.detach() | |
| G = [] | |
| for t in TASKS: | |
| g = torch.autograd.grad(branches[t], shared, retain_graph=True, allow_unused=True) | |
| G.append(torch.cat([(torch.zeros_like(p) if gi is None else gi).reshape(-1) for gi, p in zip(g, shared)]).float()) | |
| total.backward() | |
| combined = pcgrad_combine(G, self.gen) | |
| off = 0 | |
| for p in shared: | |
| k = p.numel() | |
| p.grad = combined[off:off + k].view_as(p).to(p.dtype).clone() | |
| off += k | |
| return total.detach() | |
| def famo_update(self, branches_after): | |
| """FAMO Algorithm 1, the xi step: delta_i = log D_i(theta_t) - log D_i(theta_{t+1}) on the SAME batch (branches_after: no-grad losses after opt.step()).""" | |
| if self.kind != "famo" or self._famo_prev is None: | |
| return | |
| D1 = torch.stack([(branches_after[t].detach() - self.shift).clamp_min(1e-3) for t in TASKS]) | |
| delta = torch.log(self._famo_prev) - torch.log(D1) | |
| z = torch.softmax(self.xi, -1) | |
| d = torch.autograd.grad(z, self.xi, grad_outputs=delta.detach())[0] | |
| self.xi_opt.zero_grad() | |
| self.xi.grad = d | |
| self.xi_opt.step() | |
| self._famo_prev = None | |
| def finite(self): | |
| vals = list(self.last_weights) | |
| if self.kind == "uw": | |
| vals += self.log_var.detach().cpu().tolist() | |
| if self.kind == "famo": | |
| vals += self.xi.detach().cpu().tolist() | |
| return all(math.isfinite(v) for v in vals) | |
| def pcgrad_combine(grads, gen=None, eps=1e-12): | |
| """Yu et al. 2020: grads = list of flat per-task gradient vectors (same length). For each task i, in a random order over the other tasks j, | |
| project g_i^PC away from g_j whenever g_i^PC . g_j < 0 (against the ORIGINAL g_j, as in the paper). -> the sum of the projected gradients. | |
| Orthogonal or non-conflicting task gradients pass through unchanged (identity).""" | |
| n = len(grads) | |
| pc = [g.clone() for g in grads] | |
| sq = [float(g.dot(g)) for g in grads] | |
| for i in range(n): | |
| others = [j for j in range(n) if j != i] | |
| if gen is not None and len(others) > 1: | |
| perm = torch.randperm(len(others), generator=gen).tolist() | |
| others = [others[k] for k in perm] | |
| for j in others: | |
| dot = pc[i].dot(grads[j]) | |
| if float(dot) < 0 and sq[j] > eps: | |
| pc[i] = pc[i] - (dot / sq[j]) * grads[j] | |
| return torch.stack(pc).sum(0) | |
Xet Storage Details
- Size:
- 11.8 kB
- Xet hash:
- e2cc831b25b16103ba78dfe9eb3db6e2234d3926fa99083ed71a7ba2d82df291
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.