gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: order diversity 1/3→3/3 — three root causes found and fixed

Browse Source 
1. expected_q_kernel layout mismatch: cuBLAS writes branch-major logits
   [Branch0: B×b0×NA | Branch1: ...] but kernel read sample-major
   [Sample0: 12×NA | Sample1: ...]. Fixed to use branch_logit_offset
   accumulator matching cuBLAS output layout.

2. Duplicate expected_q_kernel.cu deleted — single implementation in
   experience_kernels.cu. Trainer loads from experience_kernels.cubin.

3. Monitoring num_actions=7 (old 7-level exposure) → 9 (dir×mag = 3×3).
   Actions with dir*mag≥7 silently dropped from order/urgency counts,
   making it appear only 1 order type was active.

Result: Action diversity 15/81 (18.5%) → 45/81 (55.6%), order=3/3.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-11 23:34:53 +02:00
parent ad5c2b7ded
commit e2bf655be2
5 changed files with 43 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml/build.rs
View File

@@ -74,7 +74,6 @@ fn main() {
"relu_mask_kernel.cu",
"c51_grad_kernel.cu",
"mse_grad_kernel.cu",
"expected_q_kernel.cu",
"q_stats_kernel.cu",
"cql_grad_kernel.cu",
"trade_stats_kernel.cu",

132
crates/ml/src/cuda_pipeline/expected_q_kernel.cu
View File

@@ -1,132 +0,0 @@
/**
* Expected Q-value kernel for ad-hoc validation forward pass.
*
* Converts C51 distributional logits (value + advantage, branching dueling)
* to expected Q-values. One thread per sample; iterates over branches and atoms.
*
* Launch config: grid=(ceil(N/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void compute_expected_q(
const float* __restrict__ v_logits, // [N, num_atoms] f32
const float* __restrict__ b_logits, // [N, (b0+b1+b2+b3)*num_atoms] f32
float* __restrict__ q_values, // [N, b0+b1+b2+b3]
int N, int num_atoms,
int b0_size, int b1_size, int b2_size, int b3_size,
float v_min, float v_max)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
int total_actions = b0_size + b1_size + b2_size + b3_size;
float dz = (num_atoms > 1)
? (v_max - v_min) / (float)(num_atoms - 1)
: 0.0f;
// Value logits for this sample: [num_atoms]
const float* val = v_logits + (long long)i * num_atoms;
int branch_sizes[4];
branch_sizes[0] = b0_size;
branch_sizes[1] = b1_size;
branch_sizes[2] = b2_size;
branch_sizes[3] = b3_size;
int adv_offset = 0;
int q_offset = 0;
for (int d = 0; d < 4; d++) {
int bd = branch_sizes[d];
for (int a = 0; a < bd; a++) {
const float* adv = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + a) * num_atoms;
float eq;
if (d == 1) {
/* Magnitude (d==1) ONLY: MEAN-LOGIT Q (variance-neutral).
*
* C51's distributional softmax structurally favors zero/low-variance actions
* (Flat for direction, Small/Quarter for magnitude). Mean-logit averages
* (V + centered_A) across atoms WITHOUT softmax weighting — only the average
* advantage level matters, not the distributional shape.
*
* SAFE HERE: compute_expected_q is used for action selection + backtest eval,
* NOT for training loss. The loss kernels (mse_loss, c51_loss) have their own
* internal expected Q using proper softmax. No train/eval mismatch. */
eq = 0.0f;
for (int j = 0; j < num_atoms; j++) {
float a_mean_j = 0.0f;
for (int aa = 0; aa < bd; aa++) {
const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + aa) * num_atoms;
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float centered = adv[j] - a_mean_j;
if (d == 1) {
float sq_sum = 0.0f;
for (int aa = 0; aa < bd; aa++) {
const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + aa) * num_atoms;
float diff = adv_aa[j] - a_mean_j;
sq_sum += diff * diff;
}
float a_std = sqrtf(sq_sum / (float)bd + 1e-12f);
centered /= (a_std + 1e-6f);
}
eq += val[j] + centered;
}
eq /= (float)num_atoms;
} else {
/* Order (d==2) + urgency (d==3): standard distributional expected Q.
* C51 gradient is active for these branches — softmax is appropriate. */
float max_logit = -1e30f;
for (int j = 0; j < num_atoms; j++) {
float a_mean_j = 0.0f;
for (int aa = 0; aa < bd; aa++) {
const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + aa) * num_atoms;
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float centered = adv[j] - a_mean_j;
float combined = val[j] + centered;
max_logit = fmaxf(max_logit, combined);
}
float sum_exp = 0.0f;
for (int j = 0; j < num_atoms; j++) {
float a_mean_j = 0.0f;
for (int aa = 0; aa < bd; aa++) {
const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + aa) * num_atoms;
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float centered = adv[j] - a_mean_j;
float combined = val[j] + centered;
sum_exp += expf(combined - max_logit);
}
float log_sum = logf(sum_exp + 1e-8f) + max_logit;
eq = 0.0f;
for (int j = 0; j < num_atoms; j++) {
float a_mean_j = 0.0f;
for (int aa = 0; aa < bd; aa++) {
const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
+ (long long)(adv_offset + aa) * num_atoms;
a_mean_j += adv_aa[j];
}
a_mean_j /= (float)bd;
float centered = adv[j] - a_mean_j;
float combined = val[j] + centered;
float p = expf(combined - log_sum);
float z = v_min + (float)j * dz;
eq += p * z;
}
}
q_values[(long long)i * total_actions + q_offset + a] = bf16(eq);
}
adv_offset += bd;
q_offset += bd;
}
}

73
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -2114,44 +2114,51 @@ extern "C" __global__ void compute_expected_q(
int total_actions = b0_size + b1_size + b2_size + b3_size;
float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
/* Per-sample value logits */
/* Per-sample value logits — sample-major [N, num_atoms] */
const float* v_row = v_logits + (long long)i * num_atoms;
const float* b_row = b_logits + (long long)i * total_actions * num_atoms;
/* Compute advantage mean across all actions for dueling subtraction */
/* For each action a: dueling_logits[a][z] = v[z] + adv[a][z] - mean_adv[z]
* Expected Q = sum_z( softmax(dueling_logits[a])[z] * z_support[z] )
*
* For efficiency, compute mean advantage per atom first. */
/* Branch logits layout is BRANCH-MAJOR (cuBLAS writes per-branch blocks):
* [Branch0: N×b0×NA | Branch1: N×b1×NA | Branch2: N×b2×NA | Branch3: N×b3×NA]
* Sample i, branch d, local action a, atom z is at:
* b_logits + sum(N × n_k × NA for k < d) + i × n_d × NA + a × NA + z
*/
int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
int branch_logit_offset = 0; /* byte offset to current branch block */
int action_idx = 0; /* global action index into q_values */
for (int a = 0; a < total_actions; a++) {
const float* adv_a = b_row + (long long)a * num_atoms;
for (int d = 0; d < 4; d++) {
int n_d = branch_sizes[d];
const float* branch_base = b_logits + (long long)branch_logit_offset
+ (long long)i * n_d * num_atoms;
/* Combine value + advantage (no mean subtraction for action selection —
* the argmax is invariant to the shared value baseline). */
/* Softmax + expectation over atoms — float accumulation for numerical
* stability (softmax is sensitive to precision). */
float max_logit = -1e30f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
if (logit > max_logit) max_logit = logit;
for (int a = 0; a < n_d; a++) {
const float* adv_a = branch_base + (long long)a * num_atoms;
/* Combine value + advantage. Softmax + expectation over atoms. */
float max_logit = -1e30f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
if (logit > max_logit) max_logit = logit;
}
float sum_exp = 0.0f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
sum_exp += expf(logit - max_logit);
}
float expected_q = 0.0f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
float prob = expf(logit - max_logit) / sum_exp;
float z_val = v_min + (float)z * dz;
expected_q += prob * z_val;
}
q_values[(long long)i * total_actions + action_idx] = (expected_q);
action_idx++;
}
float sum_exp = 0.0f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
sum_exp += expf(logit - max_logit);
}
float expected_q = 0.0f;
for (int z = 0; z < num_atoms; z++) {
float logit = (v_row[z]) + (adv_a[z]);
float prob = expf(logit - max_logit) / sum_exp;
float z_val = v_min + (float)z * dz;
expected_q += prob * z_val;
}
q_values[(long long)i * total_actions + a] = (expected_q);
branch_logit_offset += N * n_d * num_atoms;
}
}

3
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -66,7 +66,8 @@ static C51_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_los
static C51_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_grad_kernel.cubin"));
static MSE_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mse_loss_kernel.cubin"));
static MSE_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mse_grad_kernel.cubin"));
static EXPECTED_Q_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/expected_q_kernel.cubin"));
/// Expected Q uses the compute_expected_q function from experience_kernels (single source of truth).
static EXPECTED_Q_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/experience_kernels.cubin"));
static Q_STATS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_stats_kernel.cubin"));
static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin"));
static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin"));

2
crates/ml/src/cuda_pipeline/gpu_monitoring.rs
View File

@@ -81,7 +81,7 @@ impl GpuMonitoringReducer {
};
// Safety: rewards and actions are valid GPU slices with at least n elements.
// summary_buf is pre-allocated 12-float device buffer on the same stream.
let num_actions_i32 = 7_i32; // exposure levels (7: ShortSmall..LongFull)
let num_actions_i32 = 9_i32; // dir(3) × mag(3) = 9 combos (was 7 from pre-4-branch era)
let order_actions_i32 = 3_i32; // DQN_ORDER_ACTIONS
let urgency_actions_i32 = 3_i32; // DQN_URGENCY_ACTIONS
unsafe {