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.
This commit is contained in:
jgrusewski
2026-04-29 16:12:57 +02:00
parent e445d07a13
commit 2cca2edbc6
2 changed files with 42 additions and 1 deletions

View File

@@ -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];

View File

@@ -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