diag+docs(dqn): trunk-slice grad decomposition + stale IQN-trunk doc fix
Extends grad_decomp_kernel to snapshot the trunk tensor slice (tensors
0..4 = w_s1, b_s1, w_s2, b_s2) in addition to the existing direction +
magnitude branch slices (8..12 / 12..16). Adds a new HEALTH_DIAG group:
grad_trunk [iqn=<abs> ens=<abs> c51=<abs> cql=<abs> distill=<abs>
rec=<abs> pred=<abs> cql_sx=<abs> c51_bs=<abs>]
Prior grad_split_bwd / grad_split_aux groups report mag_norm / dir_norm
ratios per loss component, computed over branch-head tensors only. That
measurement range structurally reports 0.0000 for any loss component
that writes exclusively to the trunk — IQN and Ens in particular. This
caused the persistent misdiagnosis that IQN-to-trunk was not wired; the
prior scoping in /tmp/foxhunt_research/iqn-to-trunk-wiring-scoping.md
confirmed the wiring is live (apply_iqn_trunk_gradient at
gpu_dqn_trainer.rs:4882) and that the zero reading was a blind spot in
the measurement pipeline.
Smoke confirms the diagnostic: after iqn_readiness ramps up (late
epochs), grad_trunk reports iqn=100..381 (real trunk SAXPY amplitude),
ens=0.07..3.57, c51=2.46..8.91 (value-head dueling path contributes
through trunk), while cql/cql_sx/distill/rec/pred stay near-zero — a
clean diagnostic baseline.
Also fixes stale documentation at dual-distributional-c51-iqn-design.md
that claimed "IQN trains in isolation — its gradients don't flow back
to the shared trunk": reworded to reflect current wired state with
file:function citation and explicit iqn_readiness gating note. Updated
the "What Changes" table ("IQN training") and "Implementation Order"
(Phase 1 marked DONE) with the same citation.
Changes:
- grad_decomp_kernel.cu: per-component result slot 2 → 3 floats
(mag_norm, dir_norm, trunk_norm); extra __shared__ sum_trunk +
tree reduction; new grad_trunk_start/trunk_len kernel args.
- gpu_dqn_trainer.rs: pinned result buffer 18 → 27 floats; snapshot
now does two copy_f32 passes (trunk → dst[0..trunk_len), branch →
dst[trunk_len..]); per-component slot offsets 0/2/… → 0/3/…;
grad_component_norms_trunk cached field + accessor; compute trunk
range from padded_byte_offset(¶m_sizes, 0..4).
- fused_training.rs: grad_trunk_norms_by_component() + per-component
grad_trunk_*_abs() accessors.
- training_loop.rs: HEALTH_DIAG emits new grad_trunk group ordered
[iqn ens c51 cql distill rec pred cql_sx c51_bs]; extended doc
comment explaining the three groups' roles.
- design spec (Problem #1 + What Changes row + Implementation Order):
stale "IQN trains in isolation" replaced by current wired-state
description, cites gpu_dqn_trainer.rs:4882 and readiness ramp at
gpu_dqn_trainer.rs:4228-4243.
Pure diagnostic — no training dynamics change, no atomicAdd, no tuning
knobs, no TF32 changes. Smokes unaffected (magnitude_distribution H10
regression pre-exists on HEAD 810b3c570; 4 other smokes pass).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -3,7 +3,17 @@
|
||||
## Problem
|
||||
|
||||
The system has both C51 (101 atoms) and IQN (32 quantiles) implemented, but:
|
||||
1. IQN trains in isolation — its gradients don't flow back to the shared trunk
|
||||
1. **[Phase 1 DONE]** IQN's backward gradient is now wired into the shared
|
||||
trunk via `apply_iqn_trunk_gradient` (`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4882-5034`),
|
||||
captured unconditionally inside the aux child graph from
|
||||
`crates/ml/src/trainers/dqn/fused_training.rs:1582-1587`. The effective
|
||||
scale is `iqn_lambda * iqn_readiness * iqn_budget`, where
|
||||
`iqn_readiness` is a loss-improvement-driven ramp (0→1) at
|
||||
`gpu_dqn_trainer.rs:4228-4243`. Trunk contribution is measurable via
|
||||
the Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group — the previous
|
||||
`grad_split_bwd [iqn=0.0000]` reading was a measurement blind spot in
|
||||
the branch-only grad_decomp snapshot range, not evidence of missing
|
||||
wiring.
|
||||
2. IQN's quantile output is never used for action selection or risk management
|
||||
3. `use_cvar_action_selection` exists in config but is hardcoded to `false`
|
||||
4. The trunk learns only from C51's cross-entropy loss, missing the risk signal from IQN
|
||||
@@ -56,7 +66,7 @@ Wire the existing IQN dual-head for two purposes:
|
||||
| Component | Today | After | Effort |
|
||||
|-----------|-------|-------|--------|
|
||||
| C51 training | Working | No change | 0 |
|
||||
| IQN training | Working, disconnected | Gradient flows to trunk | LOW |
|
||||
| IQN training | Working, trunk gradient wired (Phase 1 landed — `apply_iqn_trunk_gradient`) | No further change | DONE |
|
||||
| IQN inference | Kernel compiled, never called | Wired for CVaR | MEDIUM |
|
||||
| CVaR computation | Not implemented | New trivial CUDA kernel | LOW |
|
||||
| C51/IQN disagreement | Not implemented | |E_c51 - E_iqn| | LOW |
|
||||
@@ -67,7 +77,11 @@ Wire the existing IQN dual-head for two purposes:
|
||||
|
||||
## Implementation Order
|
||||
|
||||
1. Wire IQN gradient to shared trunk (Phase 1)
|
||||
1. ~~Wire IQN gradient to shared trunk (Phase 1)~~ — **DONE**
|
||||
(`apply_iqn_trunk_gradient` at `gpu_dqn_trainer.rs:4882-5034`,
|
||||
scheduled from `fused_training.rs:1582-1587`, scale
|
||||
`iqn_lambda * iqn_readiness * iqn_budget`). Measurable via the
|
||||
Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group.
|
||||
2. Wire IQN inference + CVaR computation (Phase 2)
|
||||
3. Add disagreement uncertainty signal (Phase 3)
|
||||
4. Hyperopt integration for cvar_alpha (Phase 4)
|
||||
|
||||
Reference in New Issue
Block a user