diff --git a/docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md b/docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md index 36349d2e3..f8b5db9e5 100644 --- a/docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md +++ b/docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md @@ -257,18 +257,84 @@ The original Steps 4-9 prescribed a separate α forward + backward kernel that b Update the launcher in `gpu_dqn_trainer.rs` to pass the 4 new args (W ptr, batch_states ptr, AUX_DIR_PROB_INDEX, STATE_DIM). For test-side callers that don't have access, pass `NULL` ptrs. - [ ] **Step 7: Thread `aux_atom_shift` into `c51_loss_kernel` (`c51_loss_kernel.cu`).** - Identify the kernel's Bellman projection code (`block_bellman_project_f` and callers — see line ~1122 of c51_loss_kernel.cu). When computing `t_z = clamp(reward + gamma * z_atom, v_min, v_max)` for the target distribution, the kernel needs to know the shifted atom positions for both: - - The PREDICTED distribution's atom positions (the predicted Q^(s, a) atoms, shifted by `aux_atom_shift[s_curr, a_curr]` for direction branch) - - The TARGET distribution's atom positions (Bellman target at next-state, shifted by `aux_atom_shift[s_next, a_next_argmax]` for direction branch) - Add same 4 args + per-call invocation logic. For non-direction branches, atom shift is 0 (existing behavior). Per `pearl_no_atomicadd`: gradient computation in the same kernel pass stays block-reduce. + **Math derivation (2026-05-13)** — the per-(b, a) atom-position shift collapses to a SINGLE per-sample reward adjustment in the Bellman projection. This dramatically reduces the kernel surgery needed. + + Define: + - `Δ_online = W[a_d] * state_121[b]` — taken action's shift on online support (a_d = a0 for dir branch) + - `Δ_target = W[a*] * next_state_121[b]` — sampled-target action's shift on target distribution (a* = best_next_a) + + Online distribution support: `z_m + Δ_online`. Target distribution atoms: `z_n + Δ_target`. Bellman backup at target atom n: + ``` + Tz_n = r + γ * (z_n + Δ_target) * (1 - done) + ``` + + Project onto online support: `b_pos_n = (Tz_n - (v_min + Δ_online)) / Δz`. Expanding: + ``` + b_pos_n = (r + γ * Δ_target * (1-done) - Δ_online + γ * z_n * (1-done) - v_min) / Δz + ``` + + **Equivalent transformation**: define `effective_reward = r + γ * Δ_target * (1-done) - Δ_online`. Then: + ``` + b_pos_n = (effective_reward + γ * z_n * (1-done) - v_min) / Δz + ``` + This is the SAME arithmetic as the unshifted projection, with `reward` replaced by `effective_reward`. The `block_bellman_project_f` helper body needs NO changes; only the `reward` arg passed in changes per dir-branch sample. + + Additionally, the argmax over actions on the next state (Step b in c51_loss_kernel) must use shifted positions: + ``` + eq_per_action[a] += W[a] * next_state_121[b] (dir branch only, all a in 0..n_d) + ``` + This shifts every action's expected Q by its own atom-shift before argmax → influences which a* gets sampled. + + **Implementation**: + 1. Add 5 new kernel args after `aux_conf_at_state`: + - `const float* __restrict__ w_aux` (`[b0_size=4]` f32, NULL-safe) + - `const float* __restrict__ batch_states` (`[B, state_dim]` f32, NULL-safe) + - `const float* __restrict__ next_batch_states` (`[B, state_dim]` f32, NULL-safe) + - `int aux_dir_prob_index` (= 121) + - `int state_dim` (= 128) + 2. For the dir branch (d==0), if `w_aux != NULL && batch_states != NULL`: + - Read `state_121[i]` + `next_state_121[i]` (one read per sample, hoisted outside the per-action loop). + - In the eq_per_action computation: after `eq = ...`, `eq += w_aux[a] * next_state_121` (BEFORE writing to `eq_per_action[a]`). + - When calling `block_bellman_project_f`: pass `effective_reward = reward + gamma_eff * Δ_target * (1-done) - Δ_online` instead of `reward`. (The existing `reward_bias` and `q_correction` adjustments stay — they compose additively at the same site.) + 3. For non-dir branches (d != 0): no change. atom-shift = 0 → effective_reward = reward (bit-identical to pre-Phase-3-α). - [ ] **Step 8: Thread `aux_atom_shift` into `c51_grad_kernel` (`c51_grad_kernel.cu`).** - The C51 gradient kernel computes `d_predicted_logit` from the projection. Under atom shift, the projection of target probability mass onto predicted atoms depends on the shifted positions. Compute `dz_shift[b, a] = Σ_n dloss/dz_n_effective` per (b, a) for direction branch. Then: - - `dW[a] += Σ_b state_121[b] * dz_shift[b, a]` (per-action block tree-reduce, no atomicAdd) - - `dstate_121[b] += Σ_a W[a] * dz_shift[b, a]` (per-sample sum, accumulate into `dbatch_states[b * state_dim + 121]`) - Both write to existing trainer buffers (`dw_aux_buf` and `dbatch_states_buf` slot 121). + **Math derivation (2026-05-13)** — gradient paths from the loss through atom-shift to W + state_121: + + Loss: `L = -Σ_k p_proj[k] * log_p_online[k]` (KL cross-entropy on projected target). + Therefore `dL/dp_proj[k] = -log_p_online[k]`. + + Projection chain rule: each target atom n with bin (lower_n, upper_n): + - `dp_proj[lower_n]/dΔ_online = +p_target_n / Δz` + - `dp_proj[upper_n]/dΔ_online = -p_target_n / Δz` + - `dp_proj[lower_n]/dΔ_target = -p_target_n * γ*(1-done) / Δz` + - `dp_proj[upper_n]/dΔ_target = +p_target_n * γ*(1-done) / Δz` + + Combining: + ``` + dL/dΔ_online = (1/Δz) × Σ_n p_target_n × (log_p_online[upper_n] - log_p_online[lower_n]) + dL/dΔ_target = -(γ*(1-done)/Δz) × Σ_n p_target_n × (log_p_online[upper_n] - log_p_online[lower_n]) + = -γ*(1-done) × dL/dΔ_online + ``` + + Chain to W and state: + ``` + dL/dW[a_d] += dL/dΔ_online × state_121[b] (dir branch only, a == a_d) + dL/dW[a*] += dL/dΔ_target × next_state_121[b] (dir branch only, a == a*) + dL/dstate_121[b] += dL/dΔ_online × W[a_d] (dir branch only) + dL/dnext_state_121[b] += dL/dΔ_target × W[a*] (dir branch only) + ``` + + **Implementation**: + 1. Add same 5 new kernel args + 2 new gradient output buffers: `dW_aux` (`[b0_size=4]` f32, accumulator) and `dbatch_states` (`[B, state_dim]` f32, slot 121 only matters). + 2. Each sample's dir-branch backward computes the `(log_p_online[upper_n] - log_p_online[lower_n])` × `p_target_n` sum (block-level via existing reduce primitives — already available in c51_grad_kernel). + 3. dW per-sample contribution → existing block tree-reduce primitives accumulate per-(b, a) → atomicAdd-free per `feedback_no_atomicadd`. + 4. dstate_121[b] writes directly to `dbatch_states[b*state_dim + 121]` (one writer per sample = no race). + 5. Per-sample writes for next_state_121: a* is per-sample (sampled action), so `dbatch_states[b*state_dim + 121]` (for the current sample's next-state buffer) is one writer. + + **Critical**: Steps 7 + 8 + 11 MUST land in the same atomic commit. Step 7 forward changes the loss landscape; Step 8 backward computes the W gradient AND the dstate_121 gradient that propagates back through the aux head. Without Step 8, the aux head's gradient is silently wrong (missing dΔ_online/dnetwork_weights path). Without Step 11, dW accumulates but W never updates → no learning. Per `feedback_no_partial_refactor.md`. - [ ] **Step 9: Thread `aux_atom_shift` into `mag_concat_qdir` (`experience_kernels.cu:5505`).** This kernel computes internal Q_dir for magnitude-branch conditioning. To keep magnitude-branch conditioning consistent with action selection, apply the same atom shift inline. Add same 4 args; shift `z_val = v_min + (float)z * dz + aux_atom_shift` per atom inside the per-action loop (lines 5560-5594 area).