Adds K independent value/advantage head weight sets sharing a common DQN
trunk. Provides uncertainty estimation (Q-value variance across heads) and
diversity regularization (KL divergence between head distributions).
Architecture:
- Head 0 stays inside CUDA Graph (zero overhead for ensemble_count=1 default)
- Heads 1..K-1 run outside CUDA Graph using post-graph save_h_s2 activations
- DtoD clone of head weights at init with stream-sync; diversity grows over training
- Pairwise KL uses symmetrized Jensen–Shannon divergence for numerical stability
New files:
- ensemble_kernels.cu: two NVRTC kernels — ensemble_aggregate_kernel (mean/var
Q-values across K heads) and ensemble_diversity_kernel (hierarchical warp→block
reduction matching dqn_grad_norm_kernel pattern, no flat atomicAdd)
- compile_ensemble_kernels() function in gpu_dqn_trainer.rs
- New GpuDqnTrainer accessors: on_v_logits_buf(), tg_h_v_scratch_ptr()
FusedTrainingCtx changes:
- ensemble_extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)>
- Pre-allocated GPU buffers (logits, mean_q, var_q, diversity_loss)
- run_ensemble_step() method runs after CUDA Graph replay (EventTrackingGuard)
- Wired in run_full_step() between IQN PER step and spectral norm step
Config: ensemble_count=1 (default, zero overhead), ensemble_diversity_weight=0.01
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
28475ff3ec