diff --git a/crates/ml/build.rs b/crates/ml/build.rs index a3940a5fe..54808caa3 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -1595,6 +1595,12 @@ fn main() { // smooths into ISV[539..543) — see alpha_isv_slots.rs. Slot 547/548 // (random baseline mean/std) are read directly from ISV by index. "phase_e_kill_criteria.cu", + // Phase E.1 Task 10 (2026-05-15): Munchausen DQN target term per + // Vieillard et al. 2020. Adds an implicit KL-regularisation bonus + // m = α_m · max(τ · log π(a|s), log_clip_min) to the TD target, + // plus a soft-V bootstrap V_soft(s') = τ · log Σ exp(Q/τ). + // One thread per batch sample; log-sum-exp-stable softmax. + "phase_e_munchausen_target.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/phase_e_munchausen_target.cu b/crates/ml/src/cuda_pipeline/phase_e_munchausen_target.cu new file mode 100644 index 000000000..863878045 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/phase_e_munchausen_target.cu @@ -0,0 +1,117 @@ +// crates/ml/src/cuda_pipeline/phase_e_munchausen_target.cu +// +// Phase E.1 Task 10 (2026-05-15): Munchausen DQN target term per Vieillard +// et al. 2020 ("Munchausen Reinforcement Learning", arXiv:2007.14430). +// +// Standard DQN target: +// target = r + γ · max_a' Q_target(s', a') +// +// Munchausen DQN target (this kernel): +// π(a|s) ∝ exp(Q_online(s, a) / τ) -- softmax policy +// m = α_m · max(τ · log π(a|s), log_clip_min) +// V_soft(s') = max_a' Q_target(s', a') + τ · log Σ_a' exp((Q_target - max) / τ) +// target = r + m + γ · V_soft(s') (if not done) +// = r + m (if done) +// +// The m term is implicit KL regularisation between successive policies; the +// soft-V bootstrap replaces the hard max with a softmax-weighted average. +// Per the paper, α_m ≈ 0.9, τ ≈ 0.03 work across Atari; we expose both as +// kernel args so an ISV-driven controller can tune them per phase. +// +// log_clip_min (typ. -1.0) prevents the m term from spiking to extreme +// negative values when π(a|s) is near zero (e.g. early training when the +// chosen action wasn't argmax). The clip preserves the *direction* of the +// regularisation (m ≤ 0 since τ·log π ≤ 0) while bounding magnitude. +// +// **Numerical stability:** both softmaxes use log-sum-exp via the max-trick +// (subtract max before exp; add it back after log). This is essential at +// τ ≈ 0.03 where raw exp(Q/τ) would overflow for any Q-spread > 25 nats. +// +// **GPU contract:** one thread per batch sample; no atomicAdd +// (`feedback_no_atomicadd`), no host branches inside the per-sample loop +// (`pearl_no_host_branches_in_captured_graph`). Inputs and outputs are +// all device pointers (or mapped-pinned for the per-sample reward stream). + +#include + +extern "C" __global__ void phase_e_munchausen_target_kernel( + /* Target net Q at next state — [batch, n_actions]. Boltzmann-softmaxed + * here to produce V_soft(s'). */ + const float* __restrict__ q_next, + /* Online net Q at current state — [batch, n_actions]. Boltzmann-softmaxed + * over actions to produce π(a|s), then we read π at the taken action. */ + const float* __restrict__ q_current, + /* Action taken — [batch]. Indexes into q_current's action dim. */ + const int* __restrict__ actions, + /* Immediate reward — [batch]. */ + const float* __restrict__ rewards, + /* Terminal mask — [batch]. 0.0 = non-terminal, 1.0 = terminal. */ + const float* __restrict__ dones, + /* DQN discount factor (typ. 0.99). */ + float gamma, + /* Munchausen scale α_m (typ. 0.9). */ + float alpha_m, + /* Boltzmann temperature τ (typ. 0.03 in the paper; smaller = sharper + * policy, more aggressive KL pull). */ + float tau, + /* Lower clip for τ · log π(a|s) (typ. -1.0). Prevents spiking. */ + float log_clip_min, + /* Output Munchausen-augmented target — [batch]. Consumed by the + * existing loss kernels in place of `r + γ · max_a' Q_target`. */ + float* __restrict__ target_out, + int batch, + int n_actions +) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch) return; + + const int row_curr = i * n_actions; + const int row_next = i * n_actions; + const float inv_tau = 1.0f / tau; + + // ---- (1) Munchausen bonus: m = α_m · max(τ · log π(a|s), log_clip_min) ---- + // + // π(a|s) ∝ exp(Q_online(s, a) / τ). The log probability of the taken + // action a* is τ · log π(a*|s) = (Q[a*] − max(Q)) − τ · log Σ exp((Q−max)/τ), + // so τ · log π(a*|s) = (Q[a*] − max(Q)) · 1 − τ · log_sum_exp_term. + // + // We compute log π(a*|s) in the standard log-sum-exp-stable way: + // log π(a*|s) = Q[a*]/τ − max(Q)/τ − log Σ exp((Q − max(Q))/τ) + // τ · log π(a*|s) = Q[a*] − max(Q) − τ · log Σ exp((Q − max(Q))/τ) + float max_curr = q_current[row_curr]; + for (int a = 1; a < n_actions; ++a) { + const float qa = q_current[row_curr + a]; + if (qa > max_curr) max_curr = qa; + } + float sum_exp_curr = 0.0f; + for (int a = 0; a < n_actions; ++a) { + sum_exp_curr += expf((q_current[row_curr + a] - max_curr) * inv_tau); + } + const int a_taken = actions[i]; + const float q_taken = q_current[row_curr + a_taken]; + const float tau_log_pi_a = (q_taken - max_curr) - tau * logf(sum_exp_curr); + const float munchausen = alpha_m * fmaxf(tau_log_pi_a, log_clip_min); + + // ---- (2) Soft V-value at next state: V_soft(s') = τ · log Σ exp(Q/τ) ---- + // + // Compute again with log-sum-exp stability over q_next. + // V_soft(s') = max(Q_next) + τ · log Σ exp((Q_next − max)/τ) + // + // Skip when the transition is terminal (dones[i] == 1.0); v_next = 0. + float v_next = 0.0f; + if (dones[i] < 0.5f) { + float max_next = q_next[row_next]; + for (int a = 1; a < n_actions; ++a) { + const float qa = q_next[row_next + a]; + if (qa > max_next) max_next = qa; + } + float sum_exp_next = 0.0f; + for (int a = 0; a < n_actions; ++a) { + sum_exp_next += expf((q_next[row_next + a] - max_next) * inv_tau); + } + v_next = max_next + tau * logf(sum_exp_next); + } + + // ---- (3) Assemble Munchausen-augmented target ---- + target_out[i] = rewards[i] + munchausen + gamma * v_next; +} diff --git a/docs/isv-slots.md b/docs/isv-slots.md index ffb86eb63..cd7063213 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -566,3 +566,24 @@ Kernel reads ISV[547]/ISV[548] (random baseline mean/std, populated once by Task **Cubin:** `target/release/build/ml-*/out/phase_e_kill_criteria.cubin` (~16.8 KB). **Launcher wiring:** lands in Phase E.1 Task 11 (the DQN training loop reads `phase_e_kill_criteria_compute_kernel` from the compiled cubin and chains `apply_pearls_ad_kernel` on the same stream after it). + +## Phase E.1 Task 10 — Munchausen target kernel (2026-05-15) + +Standalone DQN target-augmentation kernel implementing Vieillard et al. 2020 (arXiv:2007.14430). Augments the standard TD target with: + +- **Munchausen bonus:** `m = α_m · max(τ · log π(a|s), log_clip_min)` — implicit KL regularisation between successive policies (the policy is `softmax(Q_online / τ)`) +- **Soft V-bootstrap:** `V_soft(s') = max(Q_next) + τ · log Σ_a' exp((Q_next − max) / τ)` — replaces the hard max with a temperature-weighted softmax average + +Final target: `r + m + γ · V_soft(s')` (or `r + m` if terminal). + +Both softmaxes use log-sum-exp via the max-trick (essential at τ ≈ 0.03 where raw `exp(Q/τ)` would overflow f32 for any Q-spread > 25 nats). + +Does NOT touch any ISV slot directly — α_m, τ, and log_clip_min are exposed as kernel args so a downstream controller (Phase E.2+) can ISV-drive them. + +**Files touched:** +- `crates/ml/src/cuda_pipeline/phase_e_munchausen_target.cu` (new) +- `crates/ml/build.rs` — kernel added to `kernels_with_common` + +**Cubin:** `target/release/build/ml-*/out/phase_e_munchausen_target.cubin` (~12.8 KB). + +**Wiring:** lands in Phase E.1 Task 11. The integration point is wherever the existing C51/MSE loss kernels currently consume `r + γ · V_next` — they'll consume `target_out` from this kernel instead.