From 2cca2edbc6a2671758b9ecdc5d11efec32034b27 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Wed, 29 Apr 2026 16:12:57 +0200
Subject: [PATCH] chore(dqn): restore Thompson math docblocks dropped in Plan C
 T1 port

Per-function purpose comments + parameter-contract notes were present
in the Phase 0 reference (thompson_test_kernel.cu) but dropped during
the port to experience_kernels.cu. Restoring them addresses code
review feedback. No functional change.
---
 .../src/cuda_pipeline/experience_kernels.cu   | 42 ++++++++++++++++++-
 docs/dqn-wire-up-audit.md                     |  1 +
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu
index 0340d6507..38649708b 100644
--- a/crates/ml/src/cuda_pipeline/experience_kernels.cu
+++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu
@@ -165,6 +165,19 @@ __device__ __forceinline__ int argmax_n(const float* arr, int n) {
 #define N_IQN_QUANTILES 5
 #endif
 
+/**
+ * Inverse-CDF sample from a C51 atom distribution.
+ *
+ * @param p           Probability vector of length n_atoms; sums to ~1.0.
+ * @param atom_values Atom positions (support of the distribution); length n_atoms.
+ * @param n_atoms    Number of atoms (length of `p` and `atom_values`).
+ * @param u           Uniform random draw in [0, 1).
+ * @return            The atom value at the smallest index `a` with cumulative
+ *                    mass strictly greater than `u`. If floating-point roundoff
+ *                    leaves the cumulative sum below `u` after iterating all
+ *                    atoms, returns `atom_values[n_atoms - 1]` as a numerical
+ *                    fallback (documented expected behavior, not off-by-one).
+ */
 __device__ __forceinline__ float sample_c51_inverse_cdf(
     const float* __restrict__ p,
     const float* __restrict__ atom_values,
@@ -176,9 +189,19 @@ __device__ __forceinline__ float sample_c51_inverse_cdf(
         cum += p[a];
         if (u < cum) return atom_values[a];
     }
-    return atom_values[n_atoms - 1];
+    return atom_values[n_atoms - 1];  // numerical fallback
 }
 
+/**
+ * Sample from an IQN posterior by uniform-tau interpolation over the 5 fixed
+ * quantiles produced by the IQN head.
+ *
+ * @param q  Quantile values at tau in {0.05, 0.25, 0.50, 0.75, 0.95}; length
+ *           N_IQN_QUANTILES (5).
+ * @param u  Uniform random draw in [0, 1).
+ * @return   Linear interpolation of `q` at `u`. Below tau=0.05 returns q[0];
+ *           above tau=0.95 returns q[4] (clamped tails — IQN does not extrapolate).
+ */
 __device__ __forceinline__ float sample_iqn_quantile_interp(
     const float* __restrict__ q,
     float u
@@ -195,6 +218,15 @@ __device__ __forceinline__ float sample_iqn_quantile_interp(
     return q[N_IQN_QUANTILES - 1];
 }
 
+/**
+ * Compute E[Q] from C51 atom probabilities and values.
+ *
+ * @param p           Probability vector of length n_atoms; sums to ~1.0.
+ * @param atom_values Atom positions; length n_atoms.
+ * @param n_atoms    Number of atoms.
+ * @return            Sum_a p[a] * atom_values[a] — the expected return under
+ *                    the C51 categorical posterior.
+ */
 __device__ __forceinline__ float compute_e_c51_inline(
     const float* __restrict__ p,
     const float* __restrict__ atom_values,
@@ -205,6 +237,14 @@ __device__ __forceinline__ float compute_e_c51_inline(
     return e;
 }
 
+/**
+ * Compute E[Q] from IQN quantiles as the mean of the fixed-tau quantiles.
+ *
+ * @param q  Quantile values at tau in {0.05, 0.25, 0.50, 0.75, 0.95}; length
+ *           N_IQN_QUANTILES (5).
+ * @return   (1/5) * sum(q[i]) — IQN expected return under the uniform-tau
+ *           sampling convention used elsewhere in the pipeline.
+ */
 __device__ __forceinline__ float compute_e_iqn_inline(const float* __restrict__ q) {
     float s = 0.0f;
     for (int i = 0; i < N_IQN_QUANTILES; i++) s += q[i];
diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md
index e54c56cce..8356c513e 100644
--- a/docs/dqn-wire-up-audit.md
+++ b/docs/dqn-wire-up-audit.md
@@ -969,6 +969,7 @@ and reserve the un-suffixed names for any future non-inline callers.
 No callers in this commit — Task 2 wires the direction-branch Thompson
 selector inside `experience_action_select` that consumes them. Plan:
 `docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md`.
+Comments expanded post-review to preserve reference parameter-contract docs.
 
 Persistent picked_action_history scatter (2026-04-26): `gpu_backtest_evaluator.rs`
 gains a window-major `picked_action_history_buf` (parallel layout to