docs(sp22): H6 Phase 3 runbook — atom-shift math derivation for Steps 7+8

Derived the projection-arithmetic transformation that collapses per-(b, a)
atom-position shift into a SINGLE per-sample reward adjustment:

  effective_reward = reward + gamma_eff * Delta_target * (1-done) - Delta_online

where
  Delta_online = W[a_d] * state_121[b]      (taken-action shift on online support)
  Delta_target = W[a*] * next_state_121[b]  (sampled-target shift on target dist)

The Bellman projection (block_bellman_project_f body) needs NO arithmetic
change — only the reward arg passed in. This dramatically reduces Step 7's
kernel surgery surface area.

For Step 8 backward, derived the gradient paths:
  dL/dDelta_online = (1/dz) * sum_n p_target_n * (log_p_online[upper_n] - log_p_online[lower_n])
  dL/dDelta_target = -gamma*(1-done) * dL/dDelta_online

  dL/dW[a_d]            += dL/dDelta_online * state_121
  dL/dW[a*]             += dL/dDelta_target * next_state_121
  dL/dstate_121[b]      += dL/dDelta_online * W[a_d]
  dL/dnext_state_121[b] += dL/dDelta_target * W[a*]

Critical constraint documented: Steps 7+8+11 MUST land atomically. Splitting
Step 7 forward from Step 8 backward creates a silent gradient path through
the aux head (missing dDelta_online/dnetwork_weights). Splitting Step 8 from
Step 11 Adam accumulates dW without updating W — no learning. Per
feedback_no_partial_refactor.md.

The math now makes Step 7+8 a transcription job rather than exploration.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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).