fix: GPU training silently scores every sampled trajectory 0 — CPU sampling Generator vs CUDA multinomial

[kai392]

Root cause

optim.sampling.trajectory_sampling_rng (added in #171 to make --seed reproduce a run) builds a CPU torch.Generator. optim.fitness.evaluate_candidate passes it into head.select(..., sample=True), where torch.multinomial draws from softmax probabilities living on the head's device — cuda:0 per configs/trinity.yaml and CoordinatorPolicy.build's default. PyTorch requires multinomial's input and generator to share a device, so on any GPU run the first sampled turn of every trajectory raises:

RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

evaluate_candidate's asyncio.gather(..., return_exceptions=True) then swallows the error and scores the trajectory 0.0 (by design, for transient API failures), printing only [warn] N/N trajectories failed. Net effect: every candidate in every generation gets fitness 0, sep-CMA-ES optimizes a flat all-zero landscape, and the run 'completes successfully', writing a meaningless best_theta.npy. train.py always passes sample=True, run_seed=args.seed (default 0), so every real GPU training run since #171 is affected. The CPU-only CI can never see it — CPU probs + CPU generator match.

Fix

LinearHead.select now draws on the generator's device: the two prob vectors ((n_models,) / (n_roles,) — a few floats) are moved to rng.device before multinomial. rng=None behavior is unchanged (samples on the probs' own device).

Bonus: the seeded draw becomes device-independent — the same --seed picks the same (agent, role) sequence whether the head runs on CPU or CUDA, strengthening #171's reproducibility goal.

Impact / risk

• Un-breaks GPU training end-to-end; no behavior change for CPU runs (.to() is a no-op on matching devices) and none for eval (sample=False path untouched).
• Regression tests simulate the device mismatch without a GPU (meta-device generator stand-in + multinomial spy): they fail on pre-fix code, pass post-fix. The rng=None no-move contract is pinned too.

Note on CI: current main fails collection repo-wide (unresolved merge-conflict markers in src/trinity/analysis/__init__.py, issue #404 / PRs #402, #405). This PR's checks will go green once that lands; the fix itself is ruff-clean and its test file passes (21/21) locally.

🤖 Generated with Claude Code