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:
@@ -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];
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user