From 3ac51679a1472e723d4f2eb34920ef194dee05f0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 15 Mar 2026 23:50:29 +0100 Subject: [PATCH] feat(cuda): fused DQN training kernel + trainer split MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace 7k-line monolithic trainer.rs with modular trainer/ directory: action.rs, constructor.rs, metrics.rs, mod.rs, state.rs, tests.rs, training_loop.rs, train_step.rs (6048 lines total) New fused CUDA training pipeline: - dqn_training_kernel.cu: single-kernel forward+loss+backward - gpu_dqn_trainer.rs: host-side fused training orchestration - fused_training.rs: Rust-side fused training integration Eliminates per-step CPU↔GPU synchronization in DQN training loop. Co-Authored-By: Claude Opus 4.6 --- .../src/cuda_pipeline/dqn_training_kernel.cu | 1281 +++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 1595 ++++ crates/ml/src/trainers/dqn/config.rs | 34 +- crates/ml/src/trainers/dqn/fused_training.rs | 443 + crates/ml/src/trainers/dqn/mod.rs | 1 + crates/ml/src/trainers/dqn/trainer.rs | 7122 ----------------- crates/ml/src/trainers/dqn/trainer/action.rs | 396 + .../src/trainers/dqn/trainer/constructor.rs | 732 ++ crates/ml/src/trainers/dqn/trainer/metrics.rs | 782 ++ crates/ml/src/trainers/dqn/trainer/mod.rs | 872 ++ crates/ml/src/trainers/dqn/trainer/state.rs | 179 + crates/ml/src/trainers/dqn/trainer/tests.rs | 609 ++ .../ml/src/trainers/dqn/trainer/train_step.rs | 631 ++ .../src/trainers/dqn/trainer/training_loop.rs | 1847 +++++ 14 files changed, 9399 insertions(+), 7125 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/dqn_training_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs create mode 100644 crates/ml/src/trainers/dqn/fused_training.rs delete mode 100644 crates/ml/src/trainers/dqn/trainer.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/action.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/constructor.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/metrics.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/mod.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/state.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/tests.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/train_step.rs create mode 100644 crates/ml/src/trainers/dqn/trainer/training_loop.rs diff --git a/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu new file mode 100644 index 000000000..bcf1c3e90 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu @@ -0,0 +1,1281 @@ +/** + * Fused DQN training kernel — forward + C51 distributional loss in a single launch. + * + * Replaces ~160 Candle kernel dispatches (3 forward passes + C51 loss) with 1 kernel. + * Designed for CUDA Graph capture: fixed tensor shapes, no dynamic allocations. + * + * Requires common_device_functions.cuh prepended via NVRTC source concatenation. + * Launch config: grid=(batch_size, 1, 1), block=(32, 1, 1). + * One warp per sample. All 32 lanes cooperate on matrix-vector products. + * + * Network: Branching Dueling DQN with C51 distributional heads. + * - Shared layers: state → [SHARED_H1] → [SHARED_H2] + * - Value head: [SHARED_H2] → [VALUE_H] → [NUM_ATOMS] + * - Branch heads (3): [SHARED_H2] → [ADV_H] → [n_d × NUM_ATOMS] + * + * Dimension overrides injected via NVRTC #define: + * STATE_DIM, SHARED_H1, SHARED_H2, VALUE_H, ADV_H, NUM_ATOMS, + * BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE, BATCH_SIZE + */ + +/* ── Compile-time defaults (overridden by NVRTC injection) ──────────── */ +#ifndef NUM_ATOMS +#define NUM_ATOMS 51 +#endif +#ifndef BRANCH_0_SIZE +#define BRANCH_0_SIZE 5 +#endif +#ifndef BRANCH_1_SIZE +#define BRANCH_1_SIZE 3 +#endif +#ifndef BRANCH_2_SIZE +#define BRANCH_2_SIZE 3 +#endif +#ifndef BATCH_SIZE +#define BATCH_SIZE 256 +#endif +#ifndef LEAKY_RELU_ALPHA +#define LEAKY_RELU_ALPHA 0.01f +#endif + +/* Total branch output atoms (distributional) */ +#define BRANCH_0_ATOMS (BRANCH_0_SIZE * NUM_ATOMS) +#define BRANCH_1_ATOMS (BRANCH_1_SIZE * NUM_ATOMS) +#define BRANCH_2_ATOMS (BRANCH_2_SIZE * NUM_ATOMS) +#define MAX_BRANCH_SIZE 5 /* max(BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE) */ +#define NUM_BRANCHES 3 + +/* Delta between adjacent support atoms */ +#ifndef V_MIN +#define V_MIN (-25.0f) +#endif +#ifndef V_MAX +#define V_MAX (25.0f) +#endif +#define DELTA_Z ((V_MAX - V_MIN) / (float)(NUM_ATOMS - 1)) + +/* ── Shared memory size guard ───────────────────────────────────────── */ +#ifndef SHMEM_MIN +#define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b)) +#endif +#ifndef DIST_SIZE +#define DIST_SIZE(dim) (((dim) + 31) / 32) +#endif + +/* ── Device helper: warp-cooperative log_softmax ────────────────────── */ + +/** + * Compute log_softmax over `n` contiguous values starting at `logits`. + * Result written to `out`. All 32 lanes cooperate. + * Uses max-subtraction trick for numerical stability. + */ +__device__ void warp_log_softmax( + const float* logits, float* out, int n, int lane_id +) { + /* Pass 1: find max */ + float local_max = -1e30f; + for (int i = lane_id; i < n; i += 32) + local_max = fmaxf(local_max, logits[i]); + /* Warp reduce max */ + for (int offset = 16; offset > 0; offset >>= 1) + local_max = fmaxf(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, offset)); + /* Now local_max is the same in all lanes */ + + /* Pass 2: compute sum(exp(x - max)) */ + float local_sum = 0.0f; + for (int i = lane_id; i < n; i += 32) + local_sum += expf(logits[i] - local_max); + /* Warp reduce sum */ + for (int offset = 16; offset > 0; offset >>= 1) + local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, offset); + float log_sum = logf(local_sum + 1e-10f); + + /* Pass 3: write log_softmax = x - max - log(sum) */ + for (int i = lane_id; i < n; i += 32) + out[i] = logits[i] - local_max - log_sum; +} + +/** + * Compute expected value: E[X] = sum(softmax(logits) * support_atoms). + * Uses cached support atoms from shared memory to avoid recomputation. + * Falls back to inline computation if no cache provided. + */ +__device__ float warp_expected_q( + const float* log_probs, int lane_id, + const float* shmem_support /* [NUM_ATOMS] cached z_j values, or NULL */ +) { + float local_sum = 0.0f; + for (int j = lane_id; j < NUM_ATOMS; j += 32) { + float z_j = shmem_support ? shmem_support[j] : (V_MIN + j * DELTA_Z); + local_sum += expf(log_probs[j]) * z_j; + } + for (int offset = 16; offset > 0; offset >>= 1) + local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, offset); + return local_sum; +} + +/* ── Device helper: Bellman projection ──────────────────────────────── */ + +/** + * C51 Bellman projection: T_z = r + γ * z * (1-done), project onto support. + * + * Computes projected target distribution via register-local accumulation + * followed by coalesced warp-cooperative writeback, replacing the previous + * atomicAdd scatter pattern that caused non-coalesced shared memory writes. + * + * Each lane accumulates into a thread-local register array, then the warp + * reduces across lanes using __shfl_xor_sync to produce the final result. + * + * projected must be zero-initialized before this call. + * shmem_support: cached z_j values in shared memory (NULL = compute inline). + */ +__device__ void warp_bellman_project( + const float* target_probs, /* [NUM_ATOMS] target distribution */ + float reward, float done, float gamma, + float* projected, /* [NUM_ATOMS] output, must be zero-init */ + int lane_id, + const float* shmem_support /* [NUM_ATOMS] cached z_j values, or NULL */ +) { + /* ── Phase 1: Each lane accumulates projections into thread-local registers ── */ + /* Register-local bins: each lane builds its own partial projection histogram. + * This replaces the atomicAdd scatter into shared memory, eliminating + * non-coalesced writes and shared memory bank conflicts entirely. */ + float local_proj[NUM_ATOMS]; + for (int k = 0; k < NUM_ATOMS; k++) local_proj[k] = 0.0f; + + /* Each lane processes a subset of source atoms */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) { + float z_j = shmem_support ? shmem_support[j] : (V_MIN + j * DELTA_Z); + float t_z = reward + gamma * z_j * (1.0f - done); + + /* Clip to support range */ + t_z = fminf(fmaxf(t_z, V_MIN), V_MAX); + + /* Continuous index into support */ + float b = (t_z - V_MIN) / DELTA_Z; + int lower = (int)floorf(b); + int upper = (int)ceilf(b); + lower = max(lower, 0); + lower = min(lower, NUM_ATOMS - 1); + upper = max(upper, 0); + upper = min(upper, NUM_ATOMS - 1); + + float frac = b - floorf(b); + float p_j = target_probs[j]; + + /* Accumulate in registers — zero bank conflicts, zero atomics */ + local_proj[lower] += p_j * (1.0f - frac); + if (upper != lower) { + local_proj[upper] += p_j * frac; + } + } + + /* ── Phase 2: Warp-reduce each bin across all 32 lanes ── */ + /* Each bin had at most 2 contributing lanes (ceil(51/32) = 2 iterations), + * so most local_proj[k] are zero. The warp butterfly reduction merges + * the sparse contributions from all lanes into lane 0, then broadcasts. */ + for (int k = 0; k < NUM_ATOMS; k++) { + float val = local_proj[k]; + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + local_proj[k] = val; /* All lanes now have identical total */ + } + + /* ── Phase 3: Coalesced writeback to shared memory ── */ + /* Adjacent lanes write adjacent elements — perfect coalescence. */ + for (int k = lane_id; k < NUM_ATOMS; k += 32) + projected[k] = local_proj[k]; + __syncwarp(0xFFFFFFFF); +} + +/* ── Device helper: warp-cooperative forward for a single network ───── */ + +/** + * Full branching dueling forward pass with distributional (C51) output. + * + * Computes: + * shared layers → value head → branch heads → dueling → log_softmax + * + * Outputs per branch: + * - expected_q[d]: scalar expected Q per action (for argmax) + * - log_probs[d]: [n_d, NUM_ATOMS] log-softmax (for loss) + * + * For training mode (save_activations=true), saves intermediate activations + * to global memory for the backward kernel. + * + * Weight arguments use the same order as DuelingWeightSet + BranchingWeightSet. + */ +__device__ void branching_forward_distributional( + /* Input */ + const float* state_dist, + /* Shared layer weights (BF16 storage, F32 accumulation) */ + const __nv_bfloat16* __restrict__ w_s1, const __nv_bfloat16* __restrict__ b_s1, + const __nv_bfloat16* __restrict__ w_s2, const __nv_bfloat16* __restrict__ b_s2, + /* Value head weights (BF16) */ + const __nv_bfloat16* __restrict__ w_v1, const __nv_bfloat16* __restrict__ b_v1, + const __nv_bfloat16* __restrict__ w_v2, const __nv_bfloat16* __restrict__ b_v2, + /* Branch 0 (exposure) weights — from DuelingWeightSet advantage slot (BF16) */ + const __nv_bfloat16* __restrict__ w_b0_fc, const __nv_bfloat16* __restrict__ b_b0_fc, + const __nv_bfloat16* __restrict__ w_b0_out, const __nv_bfloat16* __restrict__ b_b0_out, + /* Branch 1 (order) weights (BF16) */ + const __nv_bfloat16* __restrict__ w_b1_fc, const __nv_bfloat16* __restrict__ b_b1_fc, + const __nv_bfloat16* __restrict__ w_b1_out, const __nv_bfloat16* __restrict__ b_b1_out, + /* Branch 2 (urgency) weights (BF16) */ + const __nv_bfloat16* __restrict__ w_b2_fc, const __nv_bfloat16* __restrict__ b_b2_fc, + const __nv_bfloat16* __restrict__ w_b2_out, const __nv_bfloat16* __restrict__ b_b2_out, + /* Scratch distributed buffers */ + float* scratch1_dist, + float* scratch2_dist, + /* Shared memory for weight tiles */ + float* shmem_weights, + float* shmem_bias, + /* Shared memory scratch for per-action log_softmax + Bellman projection */ + float* shmem_scratch, /* [MAX_BRANCH_SIZE * NUM_ATOMS + NUM_ATOMS] */ + /* Cached C51 support atoms in shared memory (avoids recomputation) */ + const float* shmem_support, /* [NUM_ATOMS] z_j = V_MIN + j * DELTA_Z, or NULL */ + int lane_id, + /* Outputs (register arrays, local to this thread) */ + float* branch_expected_q, /* [MAX_BRANCH_SIZE * NUM_BRANCHES] = [5+3+3=11] expected Q values */ + float* branch_log_probs, /* [BRANCH_0_ATOMS + BRANCH_1_ATOMS + BRANCH_2_ATOMS] all log-probs */ + /* Activation saves (global memory, per-sample) */ + float* save_h_s1, /* [SHARED_H1] — NULL if not saving */ + float* save_h_s2, /* [SHARED_H2] — NULL if not saving */ + float* save_h_v, /* [VALUE_H] — NULL if not saving */ + float* save_h_b0, /* [ADV_H] — NULL if not saving */ + float* save_h_b1, /* [ADV_H] */ + float* save_h_b2 /* [ADV_H] */ +) { + /* ── Shared layers (BF16 weights, F32 activations) ────────────── */ + /* state → scratch1 (h_s1) */ + TILE_LAYER_WARP_BF16(w_s1, b_s1, state_dist, scratch1_dist, + STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, lane_id); + /* scratch1 (h_s1) → scratch2 (h_s2) */ + TILE_LAYER_WARP_BF16(w_s2, b_s2, scratch1_dist, scratch2_dist, + SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, lane_id); + + /* Save h_s2 if needed (for backward pass) */ + if (save_h_s2) { + for (int i = lane_id; i < SHARED_H2; i += 32) + save_h_s2[i] = scratch2_dist[i / 32]; + } + /* Save h_s1 if needed — scratch1_dist still holds h_s1 here + * (will be overwritten for h_v below). Needed for backward through shared_1. */ + if (save_h_s1) { + for (int i = lane_id; i < SHARED_H1; i += 32) + save_h_s1[i] = scratch1_dist[i / 32]; + } + + /* ── Value head (BF16 weights) ─────────────────────────────────── */ + float* h_v_dist = scratch1_dist; /* reuse scratch1 for value hidden */ + TILE_LAYER_WARP_BF16(w_v1, b_v1, scratch2_dist, h_v_dist, + SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, lane_id); + + /* Save h_v if needed */ + if (save_h_v) { + for (int i = lane_id; i < VALUE_H; i += 32) + save_h_v[i] = h_v_dist[i / 32]; + } + + /* Value output: [NUM_ATOMS, VALUE_H] → value_logits[NUM_ATOMS] (BF16 weights) */ + float value_logits[NUM_ATOMS]; + { + __nv_bfloat16* shmem_w_bf16 = (__nv_bfloat16*)shmem_weights; + for (int tile = 0; tile < (NUM_ATOMS + SHMEM_TILE_ROWS_BF16 - 1) / SHMEM_TILE_ROWS_BF16; tile++) { + int ts = tile * SHMEM_TILE_ROWS_BF16; + int tr = SHMEM_MIN(SHMEM_TILE_ROWS_BF16, NUM_ATOMS - ts); + cooperative_load_tile_bf16(shmem_w_bf16, w_v2 + ts * VALUE_H, tr * VALUE_H); + cooperative_load_bias_bf16_to_f32(shmem_bias, b_v2 + ts, tr); + __syncwarp(0xFFFFFFFF); + for (int r = 0; r < tr; r++) { + int out_idx = ts + r; + if (out_idx >= NUM_ATOMS) break; + const __nv_bfloat16* row = shmem_w_bf16 + r * VALUE_H; + float partial = 0.0f; + for (int i = lane_id; i < VALUE_H; i += 32) + partial += __bfloat162float(row[i]) * h_v_dist[i / 32]; + if (lane_id == 0) partial += shmem_bias[r]; + value_logits[out_idx] = warp_reduce_sum_all(partial); + } + __syncwarp(0xFFFFFFFF); + } + } + /* value_logits: [NUM_ATOMS] — all lanes have identical copy */ + + /* ── Branch heads ───────────────────────────────────────────────── */ + /* Process each branch sequentially to limit register pressure. + * For each branch: FC → LeakyReLU → output → dueling → log_softmax. */ + + /* Branch weight pointers and sizes (compile-time known, BF16) */ + const __nv_bfloat16* branch_fc_w[3] = { w_b0_fc, w_b1_fc, w_b2_fc }; + const __nv_bfloat16* branch_fc_b[3] = { b_b0_fc, b_b1_fc, b_b2_fc }; + const __nv_bfloat16* branch_out_w[3] = { w_b0_out, w_b1_out, w_b2_out }; + const __nv_bfloat16* branch_out_b[3] = { b_b0_out, b_b1_out, b_b2_out }; + const int branch_sizes[3] = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE }; + const int branch_atoms[3] = { BRANCH_0_ATOMS, BRANCH_1_ATOMS, BRANCH_2_ATOMS }; + float* branch_save_h[3] = { save_h_b0, save_h_b1, save_h_b2 }; + + int lp_offset = 0; /* offset into branch_log_probs output */ + int eq_offset = 0; /* offset into branch_expected_q output */ + + for (int d = 0; d < NUM_BRANCHES; d++) { + int n_d = branch_sizes[d]; + int n_atoms = branch_atoms[d]; + + /* Branch FC: h_s2 → h_bd (BF16 weights) */ + float* h_bd_dist = scratch1_dist + DIST_SIZE(VALUE_H); /* after h_v space */ + TILE_LAYER_WARP_BF16(branch_fc_w[d], branch_fc_b[d], scratch2_dist, h_bd_dist, + SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, lane_id); + + /* Save h_bd if needed */ + if (branch_save_h[d]) { + for (int i = lane_id; i < ADV_H; i += 32) + branch_save_h[d][i] = h_bd_dist[i / 32]; + } + + /* Branch output: [n_d * NUM_ATOMS, ADV_H] → raw_logits in shmem_scratch (BF16) */ + float* raw_logits = shmem_scratch; /* [n_d * NUM_ATOMS] */ + { + __nv_bfloat16* shmem_w_bf16 = (__nv_bfloat16*)shmem_weights; + for (int tile = 0; tile < (n_atoms + SHMEM_TILE_ROWS_BF16 - 1) / SHMEM_TILE_ROWS_BF16; tile++) { + int ts = tile * SHMEM_TILE_ROWS_BF16; + int tr = SHMEM_MIN(SHMEM_TILE_ROWS_BF16, n_atoms - ts); + cooperative_load_tile_bf16(shmem_w_bf16, branch_out_w[d] + ts * ADV_H, tr * ADV_H); + cooperative_load_bias_bf16_to_f32(shmem_bias, branch_out_b[d] + ts, tr); + __syncwarp(0xFFFFFFFF); + for (int r = 0; r < tr; r++) { + int out_idx = ts + r; + if (out_idx >= n_atoms) break; + const __nv_bfloat16* row = shmem_w_bf16 + r * ADV_H; + float partial = 0.0f; + for (int i = lane_id; i < ADV_H; i += 32) + partial += __bfloat162float(row[i]) * h_bd_dist[i / 32]; + if (lane_id == 0) partial += shmem_bias[r]; + float val = warp_reduce_sum_all(partial); + if (lane_id == 0) raw_logits[out_idx] = val; + } + __syncwarp(0xFFFFFFFF); + } + } + /* raw_logits now in shared memory: [n_d * NUM_ATOMS] (lane 0 wrote them) */ + /* Broadcast to all lanes via shared memory read */ + + /* Dueling: Q[a,j] = V[j] + A[a,j] - mean_a(A[*,j]) for each atom j. + * Parallelized across warp: each lane handles different atom indices. */ + float* dueling_logits = shmem_scratch; /* overwrite in-place */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) { + float a_mean = 0.0f; + for (int a = 0; a < n_d; a++) + a_mean += raw_logits[a * NUM_ATOMS + j]; + a_mean /= (float)n_d; + for (int a = 0; a < n_d; a++) { + int idx_local = a * NUM_ATOMS + j; + dueling_logits[idx_local] = value_logits[j] + raw_logits[idx_local] - a_mean; + } + } + __syncwarp(0xFFFFFFFF); + + /* Log-softmax per action (along atoms dim) */ + for (int a = 0; a < n_d; a++) { + float* action_logits = dueling_logits + a * NUM_ATOMS; + float* action_lp = &branch_log_probs[lp_offset + a * NUM_ATOMS]; + warp_log_softmax(action_logits, action_lp, NUM_ATOMS, lane_id); + + /* Expected Q for this action (uses cached support atoms) */ + branch_expected_q[eq_offset + a] = warp_expected_q(action_lp, lane_id, shmem_support); + } + + lp_offset += n_atoms; + eq_offset += n_d; + } +} + +/* ── Main kernel: Forward + C51 Loss ────────────────────────────────── */ + +extern "C" __global__ void dqn_forward_loss_kernel( + /* ── Batch data (from GPU PER sampling) ──────────────────────── */ + const float* __restrict__ states, /* [B, STATE_DIM] */ + const float* __restrict__ next_states, /* [B, STATE_DIM] */ + const int* __restrict__ actions, /* [B] factored action indices 0-44 */ + const float* __restrict__ rewards, /* [B] */ + const float* __restrict__ dones, /* [B] */ + const float* __restrict__ is_weights, /* [B] PER importance-sampling weights */ + + /* ── Online network weights (20 BF16 tensors — tensor core throughput) ─ */ + const __nv_bfloat16* __restrict__ on_w_s1, const __nv_bfloat16* __restrict__ on_b_s1, + const __nv_bfloat16* __restrict__ on_w_s2, const __nv_bfloat16* __restrict__ on_b_s2, + const __nv_bfloat16* __restrict__ on_w_v1, const __nv_bfloat16* __restrict__ on_b_v1, + const __nv_bfloat16* __restrict__ on_w_v2, const __nv_bfloat16* __restrict__ on_b_v2, + /* Branch 0 = exposure (from DuelingWeightSet advantage slot, BF16) */ + const __nv_bfloat16* __restrict__ on_w_b0fc, const __nv_bfloat16* __restrict__ on_b_b0fc, + const __nv_bfloat16* __restrict__ on_w_b0out, const __nv_bfloat16* __restrict__ on_b_b0out, + /* Branch 1 = order (from BranchingWeightSet, BF16) */ + const __nv_bfloat16* __restrict__ on_w_b1fc, const __nv_bfloat16* __restrict__ on_b_b1fc, + const __nv_bfloat16* __restrict__ on_w_b1out, const __nv_bfloat16* __restrict__ on_b_b1out, + /* Branch 2 = urgency (from BranchingWeightSet, BF16) */ + const __nv_bfloat16* __restrict__ on_w_b2fc, const __nv_bfloat16* __restrict__ on_b_b2fc, + const __nv_bfloat16* __restrict__ on_w_b2out, const __nv_bfloat16* __restrict__ on_b_b2out, + + /* ── Target network weights (same BF16 layout) ────────────────── */ + const __nv_bfloat16* __restrict__ tg_w_s1, const __nv_bfloat16* __restrict__ tg_b_s1, + const __nv_bfloat16* __restrict__ tg_w_s2, const __nv_bfloat16* __restrict__ tg_b_s2, + const __nv_bfloat16* __restrict__ tg_w_v1, const __nv_bfloat16* __restrict__ tg_b_v1, + const __nv_bfloat16* __restrict__ tg_w_v2, const __nv_bfloat16* __restrict__ tg_b_v2, + const __nv_bfloat16* __restrict__ tg_w_b0fc, const __nv_bfloat16* __restrict__ tg_b_b0fc, + const __nv_bfloat16* __restrict__ tg_w_b0out, const __nv_bfloat16* __restrict__ tg_b_b0out, + const __nv_bfloat16* __restrict__ tg_w_b1fc, const __nv_bfloat16* __restrict__ tg_b_b1fc, + const __nv_bfloat16* __restrict__ tg_w_b1out, const __nv_bfloat16* __restrict__ tg_b_b1out, + const __nv_bfloat16* __restrict__ tg_w_b2fc, const __nv_bfloat16* __restrict__ tg_b_b2fc, + const __nv_bfloat16* __restrict__ tg_w_b2out, const __nv_bfloat16* __restrict__ tg_b_b2out, + + /* ── Saved activations (for backward kernel) ─────────────────── */ + float* __restrict__ save_h_s1, /* [B, SHARED_H1] */ + float* __restrict__ save_h_s2, /* [B, SHARED_H2] */ + float* __restrict__ save_h_v, /* [B, VALUE_H] */ + float* __restrict__ save_h_b0, /* [B, ADV_H] */ + float* __restrict__ save_h_b1, /* [B, ADV_H] */ + float* __restrict__ save_h_b2, /* [B, ADV_H] */ + float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, NUM_ATOMS] taken action log-probs */ + float* __restrict__ save_projected, /* [B, NUM_BRANCHES, NUM_ATOMS] Bellman projected targets */ + + /* ── Outputs ─────────────────────────────────────────────────── */ + float* __restrict__ out_per_sample_loss, /* [B] weighted loss per sample */ + float* __restrict__ out_td_errors, /* [B] for PER priority update */ + float* __restrict__ out_total_loss, /* [1] batch mean loss */ + + /* ── Config ──────────────────────────────────────────────────── */ + float gamma, + int batch_size +) { + extern __shared__ float shmem[]; + /* Shared memory layout (extended with state cache + C51 support cache): + * [0, SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM) — weight tile (F32 region, also holds 2× BF16 rows) + * [weight_end, weight_end + SHMEM_TILE_ROWS_BF16) — bias tile (F32, sized for doubled BF16 tiles) + * [bias_end, bias_end + MAX_BRANCH_SIZE * NUM_ATOMS + NUM_ATOMS) — scratch for logits + projection + * [scratch_end, scratch_end + STATE_DIM) — state vector cache + * [state_end, state_end + NUM_ATOMS) — C51 support atoms cache + */ + float* shmem_weights = shmem; + float* shmem_bias = shmem + SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM; + float* shmem_scratch = shmem_bias + SHMEM_TILE_ROWS_BF16; + float* shmem_state = shmem_scratch + MAX_BRANCH_SIZE * NUM_ATOMS + NUM_ATOMS; + float* shmem_support = shmem_state + STATE_DIM; + + int sample_id = blockIdx.x; + int lane_id = threadIdx.x; + if (sample_id >= batch_size) return; + + /* ── Populate C51 support cache (once per block, constant across samples) ── */ + /* z_j = V_MIN + j * DELTA_Z for j in [0, NUM_ATOMS). + * Eliminates repeated FMA per atom in warp_expected_q and warp_bellman_project. + * Only 51 floats (204 bytes) — negligible shmem cost, used ~33 times per sample + * (11 expected_q calls across 3 forward passes + 3 Bellman projections). */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) + shmem_support[j] = V_MIN + j * DELTA_Z; + __syncwarp(0xFFFFFFFF); + + /* ── Load state into shared memory cache + distributed registers ── */ + /* State is loaded from global memory ONCE into shared memory, then + * distributed into registers. The shmem copy persists for the online + * forward pass (PASS 1); next_states reuses shmem_state for PASS 2/3. */ + for (int i = lane_id; i < STATE_DIM; i += 32) + shmem_state[i] = states[sample_id * STATE_DIM + i]; + __syncwarp(0xFFFFFFFF); + + float state_dist[DIST_SIZE(STATE_DIM)]; + for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) state_dist[i] = 0.0f; + for (int i = lane_id; i < STATE_DIM; i += 32) + state_dist[i / 32] = shmem_state[i]; + + /* ── Scratch distributed buffers ────────────────────────────── */ + float scratch1_dist[DIST_SIZE(SHARED_H1)]; + float scratch2_dist[DIST_SIZE(SHARED_H2)]; + + /* ── Output buffers (register arrays) ───────────────────────── */ + /* All log-probs for all branches: BRANCH_0_ATOMS + BRANCH_1_ATOMS + BRANCH_2_ATOMS */ + float online_log_probs[BRANCH_0_ATOMS + BRANCH_1_ATOMS + BRANCH_2_ATOMS]; + float online_expected_q[BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE]; + + /* ═══════════════════════════════════════════════════════════════ + * PASS 1: Online forward on STATES (training mode, save activations) + * ═══════════════════════════════════════════════════════════════ */ + branching_forward_distributional( + state_dist, + on_w_s1, on_b_s1, on_w_s2, on_b_s2, + on_w_v1, on_b_v1, on_w_v2, on_b_v2, + on_w_b0fc, on_b_b0fc, on_w_b0out, on_b_b0out, + on_w_b1fc, on_b_b1fc, on_w_b1out, on_b_b1out, + on_w_b2fc, on_b_b2fc, on_w_b2out, on_b_b2out, + scratch1_dist, scratch2_dist, + shmem_weights, shmem_bias, shmem_scratch, + shmem_support, + lane_id, + online_expected_q, online_log_probs, + /* Save activations for backward pass */ + save_h_s1 + sample_id * SHARED_H1, + save_h_s2 + sample_id * SHARED_H2, + save_h_v + sample_id * VALUE_H, + save_h_b0 + sample_id * ADV_H, + save_h_b1 + sample_id * ADV_H, + save_h_b2 + sample_id * ADV_H + ); + + /* ═══════════════════════════════════════════════════════════════ + * PASS 2: Target forward on NEXT_STATES (inference, no saves) + * ═══════════════════════════════════════════════════════════════ */ + /* Reload shmem_state with next_states (reuse same cache slot) */ + for (int i = lane_id; i < STATE_DIM; i += 32) + shmem_state[i] = next_states[sample_id * STATE_DIM + i]; + __syncwarp(0xFFFFFFFF); + + float next_state_dist[DIST_SIZE(STATE_DIM)]; + for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) next_state_dist[i] = 0.0f; + for (int i = lane_id; i < STATE_DIM; i += 32) + next_state_dist[i / 32] = shmem_state[i]; + + float target_log_probs[BRANCH_0_ATOMS + BRANCH_1_ATOMS + BRANCH_2_ATOMS]; + float target_expected_q[BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE]; + + branching_forward_distributional( + next_state_dist, + tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, + tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, + tg_w_b0fc, tg_b_b0fc, tg_w_b0out, tg_b_b0out, + tg_w_b1fc, tg_b_b1fc, tg_w_b1out, tg_b_b1out, + tg_w_b2fc, tg_b_b2fc, tg_w_b2out, tg_b_b2out, + scratch1_dist, scratch2_dist, + shmem_weights, shmem_bias, shmem_scratch, + shmem_support, + lane_id, + target_expected_q, target_log_probs, + NULL, NULL, NULL, NULL, NULL, NULL /* no activation saves */ + ); + + /* ═══════════════════════════════════════════════════════════════ + * PASS 3: Online forward on NEXT_STATES (Double DQN action selection) + * ═══════════════════════════════════════════════════════════════ */ + float online_next_expected_q[BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE]; + /* We don't need log-probs from this pass, only expected Q for argmax. + * Reuse online_log_probs as scratch (overwritten, not needed after loss). */ + float online_next_lp_scratch[BRANCH_0_ATOMS + BRANCH_1_ATOMS + BRANCH_2_ATOMS]; + + branching_forward_distributional( + next_state_dist, + on_w_s1, on_b_s1, on_w_s2, on_b_s2, + on_w_v1, on_b_v1, on_w_v2, on_b_v2, + on_w_b0fc, on_b_b0fc, on_w_b0out, on_b_b0out, + on_w_b1fc, on_b_b1fc, on_w_b1out, on_b_b1out, + on_w_b2fc, on_b_b2fc, on_w_b2out, on_b_b2out, + scratch1_dist, scratch2_dist, + shmem_weights, shmem_bias, shmem_scratch, + shmem_support, + lane_id, + online_next_expected_q, online_next_lp_scratch, + NULL, NULL, NULL, NULL, NULL, NULL + ); + + /* ═══════════════════════════════════════════════════════════════ + * C51 DISTRIBUTIONAL LOSS (per-branch cross-entropy) + * ═══════════════════════════════════════════════════════════════ */ + + /* Decompose factored action into per-branch indices */ + int factored_action = actions[sample_id]; + int branch_action[3]; + branch_action[0] = factored_action / 9; /* exposure: 0-4 */ + branch_action[1] = (factored_action % 9) / 3; /* order: 0-2 */ + branch_action[2] = factored_action % 3; /* urgency: 0-2 */ + + float reward = rewards[sample_id]; + float done = dones[sample_id]; + float is_weight = is_weights[sample_id]; + + const int branch_sizes[3] = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE }; + const int branch_atom_counts[3] = { BRANCH_0_ATOMS, BRANCH_1_ATOMS, BRANCH_2_ATOMS }; + + float total_ce = 0.0f; /* accumulated cross-entropy over branches */ + + int lp_off = 0; /* offset into online_log_probs / target_log_probs */ + int eq_off = 0; /* offset into expected_q arrays */ + + for (int d = 0; d < NUM_BRANCHES; d++) { + int n_d = branch_sizes[d]; + int a_d = branch_action[d]; + + /* ── Step 1: Current log-probs for taken action ─────────── */ + /* online_log_probs[lp_off + a_d * NUM_ATOMS ... + NUM_ATOMS] */ + const float* current_lp = &online_log_probs[lp_off + a_d * NUM_ATOMS]; + + /* Save for backward kernel */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) + save_current_lp[sample_id * NUM_BRANCHES * NUM_ATOMS + d * NUM_ATOMS + j] = current_lp[j]; + + /* ── Step 2: Best next action (Double DQN: online selects) */ + int best_next_a = 0; + float best_q = online_next_expected_q[eq_off]; + for (int a = 1; a < n_d; a++) { + float q = online_next_expected_q[eq_off + a]; + if (q > best_q) { + best_q = q; + best_next_a = a; + } + } + + /* ── Step 3: Target probs for best next action ──────────── */ + const float* target_lp_best = &target_log_probs[lp_off + best_next_a * NUM_ATOMS]; + /* Convert log-probs to probs for Bellman projection */ + float target_probs_best[NUM_ATOMS]; + for (int j = 0; j < NUM_ATOMS; j++) + target_probs_best[j] = expf(target_lp_best[j]); + + /* ── Step 4: Bellman projection ─────────────────────────── */ + float* projected = shmem_scratch; /* reuse shared memory */ + /* Zero-initialize projected array */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) + projected[j] = 0.0f; + __syncwarp(0xFFFFFFFF); + + warp_bellman_project(target_probs_best, reward, done, gamma, projected, lane_id, shmem_support); + + /* Save projected targets for backward kernel */ + for (int j = lane_id; j < NUM_ATOMS; j += 32) + save_projected[sample_id * NUM_BRANCHES * NUM_ATOMS + d * NUM_ATOMS + j] = projected[j]; + + /* ── Step 5: Cross-entropy ──────────────────────────────── */ + /* CE_d = -sum_j(projected[j] * current_lp[j]) */ + float local_ce = 0.0f; + for (int j = lane_id; j < NUM_ATOMS; j += 32) + local_ce -= projected[j] * current_lp[j]; + /* Warp reduce */ + for (int offset = 16; offset > 0; offset >>= 1) + local_ce += __shfl_xor_sync(0xFFFFFFFF, local_ce, offset); + + total_ce += local_ce; + + lp_off += branch_atom_counts[d]; + eq_off += n_d; + } + + /* ── Average over branches, apply IS weight ─────────────────── */ + float per_sample_loss = (total_ce / (float)NUM_BRANCHES) * is_weight; + float td_error = total_ce / (float)NUM_BRANCHES; /* unweighted for PER priorities */ + + /* ── Write outputs (lane 0 only) ────────────────────────────── */ + if (lane_id == 0) { + out_per_sample_loss[sample_id] = per_sample_loss; + out_td_errors[sample_id] = td_error; + atomicAdd(out_total_loss, per_sample_loss / (float)batch_size); + } +} + +/* ══════════════════════════════════════════════════════════════════════ + * GRADIENT BUFFER LAYOUT + * + * All parameter gradients are accumulated into a single flat buffer via + * atomicAdd. This layout matches the parameter flattening order used by + * the Rust host for Adam optimizer state. + * ══════════════════════════════════════════════════════════════════════ */ + +/* Parameter sizes (element counts, not bytes) */ +#define PARAM_W_S1 (SHARED_H1 * STATE_DIM) +#define PARAM_B_S1 (SHARED_H1) +#define PARAM_W_S2 (SHARED_H2 * SHARED_H1) +#define PARAM_B_S2 (SHARED_H2) +#define PARAM_W_V1 (VALUE_H * SHARED_H2) +#define PARAM_B_V1 (VALUE_H) +#define PARAM_W_V2 (NUM_ATOMS * VALUE_H) +#define PARAM_B_V2 (NUM_ATOMS) +#define PARAM_W_B0FC (ADV_H * SHARED_H2) +#define PARAM_B_B0FC (ADV_H) +#define PARAM_W_B0OUT (BRANCH_0_ATOMS * ADV_H) +#define PARAM_B_B0OUT (BRANCH_0_ATOMS) +#define PARAM_W_B1FC (ADV_H * SHARED_H2) +#define PARAM_B_B1FC (ADV_H) +#define PARAM_W_B1OUT (BRANCH_1_ATOMS * ADV_H) +#define PARAM_B_B1OUT (BRANCH_1_ATOMS) +#define PARAM_W_B2FC (ADV_H * SHARED_H2) +#define PARAM_B_B2FC (ADV_H) +#define PARAM_W_B2OUT (BRANCH_2_ATOMS * ADV_H) +#define PARAM_B_B2OUT (BRANCH_2_ATOMS) + +/* Cumulative offsets into flat gradient buffer */ +#define GOFF_W_S1 0 +#define GOFF_B_S1 (GOFF_W_S1 + PARAM_W_S1) +#define GOFF_W_S2 (GOFF_B_S1 + PARAM_B_S1) +#define GOFF_B_S2 (GOFF_W_S2 + PARAM_W_S2) +#define GOFF_W_V1 (GOFF_B_S2 + PARAM_B_S2) +#define GOFF_B_V1 (GOFF_W_V1 + PARAM_W_V1) +#define GOFF_W_V2 (GOFF_B_V1 + PARAM_B_V1) +#define GOFF_B_V2 (GOFF_W_V2 + PARAM_W_V2) +#define GOFF_W_B0FC (GOFF_B_V2 + PARAM_B_V2) +#define GOFF_B_B0FC (GOFF_W_B0FC + PARAM_W_B0FC) +#define GOFF_W_B0OUT (GOFF_B_B0FC + PARAM_B_B0FC) +#define GOFF_B_B0OUT (GOFF_W_B0OUT + PARAM_W_B0OUT) +#define GOFF_W_B1FC (GOFF_B_B0OUT + PARAM_B_B0OUT) +#define GOFF_B_B1FC (GOFF_W_B1FC + PARAM_W_B1FC) +#define GOFF_W_B1OUT (GOFF_B_B1FC + PARAM_B_B1FC) +#define GOFF_B_B1OUT (GOFF_W_B1OUT + PARAM_W_B1OUT) +#define GOFF_W_B2FC (GOFF_B_B1OUT + PARAM_B_B1OUT) +#define GOFF_B_B2FC (GOFF_W_B2FC + PARAM_W_B2FC) +#define GOFF_W_B2OUT (GOFF_B_B2FC + PARAM_B_B2FC) +#define GOFF_B_B2OUT (GOFF_W_B2OUT + PARAM_W_B2OUT) +#define TOTAL_PARAMS (GOFF_B_B2OUT + PARAM_B_B2OUT) + +/* ══════════════════════════════════════════════════════════════════════ + * BACKWARD KERNEL + * + * Computes gradients through: + * C51 cross-entropy → log_softmax → dueling → linear layers + * Accumulates into a single flat gradient buffer via atomicAdd. + * + * Launch config: grid=(batch_size, 1, 1), block=(32, 1, 1). + * One warp per sample (same as forward kernel). + * ══════════════════════════════════════════════════════════════════════ */ + +/** + * Backward through a linear output layer (no activation). + * Tiled: processes SHMEM_TILE_ROWS rows of W at a time. + * + * For each output row i: + * grad_W[i,j] += dL_dy[i] * x[j] (atomicAdd) + * grad_b[i] += dL_dy[i] (atomicAdd, lane 0) + * dL_dx[j] += W[i,j] * dL_dy[i] (local accumulate) + * + * dL_dy[i] is computed on-the-fly from dL_dlogit and dueling factors. + */ +__device__ void backward_output_tiled( + const float* dL_dlogit, /* [NUM_ATOMS] gradient w.r.t. dueling logits */ + int a_d, /* taken action index for this branch */ + int n_d, /* number of actions for this branch */ + float scale, /* is_weight / NUM_BRANCHES */ + const float* x_dist, /* [DIST_SIZE(in_dim)] input activation (distributed) */ + const float* __restrict__ W, /* [n_d * NUM_ATOMS, in_dim] weight matrix */ + float* grad_W, /* atomicAdd target for weight grads */ + float* grad_b, /* atomicAdd target for bias grads */ + float* dL_dx_dist, /* [DIST_SIZE(in_dim)] accumulated input gradient */ + int in_dim, + int total_out, /* n_d * NUM_ATOMS */ + float* shmem_weights, + int lane_id +) { + for (int tile = 0; tile < (total_out + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; tile++) { + int ts = tile * SHMEM_TILE_ROWS; + int tr = SHMEM_MIN(SHMEM_TILE_ROWS, total_out - ts); + cooperative_load_tile(shmem_weights, W + ts * in_dim, tr * in_dim); + __syncwarp(0xFFFFFFFF); + + for (int r = 0; r < tr; r++) { + int out_idx = ts + r; + if (out_idx >= total_out) break; + + /* Compute dL/dy for this output index on-the-fly */ + int a = out_idx / NUM_ATOMS; + int k = out_idx % NUM_ATOMS; + float factor = (a == a_d) ? (1.0f - 1.0f / (float)n_d) : (-1.0f / (float)n_d); + float dy = dL_dlogit[k] * factor * scale; + + const float* W_row = shmem_weights + r * in_dim; + + /* Weight gradient: atomicAdd */ + for (int j = lane_id; j < in_dim; j += 32) + atomicAdd(&grad_W[out_idx * in_dim + j], dy * x_dist[j / 32]); + + /* Bias gradient */ + if (lane_id == 0) + atomicAdd(&grad_b[out_idx], dy); + + /* Input gradient: W^T @ dL_dy */ + for (int j = lane_id; j < in_dim; j += 32) + dL_dx_dist[j / 32] += W_row[j] * dy; + } + __syncwarp(0xFFFFFFFF); + } +} + +/** + * Backward through value output layer (no activation, no dueling factoring). + * dL_dy = dL_dvalue_logits directly (already accumulated across branches). + */ +__device__ void backward_value_out_tiled( + const float* dL_dvalue, /* [NUM_ATOMS] */ + float scale, /* is_weight / NUM_BRANCHES */ + const float* x_dist, /* [DIST_SIZE(VALUE_H)] = h_v distributed */ + const float* __restrict__ W, /* [NUM_ATOMS, VALUE_H] */ + float* grad_W, /* atomicAdd target */ + float* grad_b, /* atomicAdd target */ + float* dL_dx_dist, /* [DIST_SIZE(VALUE_H)] accumulated */ + float* shmem_weights, + int lane_id +) { + for (int tile = 0; tile < (NUM_ATOMS + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; tile++) { + int ts = tile * SHMEM_TILE_ROWS; + int tr = SHMEM_MIN(SHMEM_TILE_ROWS, NUM_ATOMS - ts); + cooperative_load_tile(shmem_weights, W + ts * VALUE_H, tr * VALUE_H); + __syncwarp(0xFFFFFFFF); + + for (int r = 0; r < tr; r++) { + int k = ts + r; + if (k >= NUM_ATOMS) break; + float dy = dL_dvalue[k] * scale; + const float* W_row = shmem_weights + r * VALUE_H; + + for (int j = lane_id; j < VALUE_H; j += 32) + atomicAdd(&grad_W[k * VALUE_H + j], dy * x_dist[j / 32]); + if (lane_id == 0) + atomicAdd(&grad_b[k], dy); + for (int j = lane_id; j < VALUE_H; j += 32) + dL_dx_dist[j / 32] += W_row[j] * dy; + } + __syncwarp(0xFFFFFFFF); + } +} + +/** + * Backward through FC + LeakyReLU layer (tiled). + * Given dL/dh (gradient of post-activation), computes: + * dL/dz = dL/dh * leaky_relu_grad(h) where h = saved activation + * grad_W[i,j] += dL_dz[i] * x[j] (atomicAdd) + * grad_b[i] += dL_dz[i] (atomicAdd, lane 0) + * dL_dx[j] += W[i,j] * dL_dz[i] (local accumulate, skipped if !compute_input_grad) + * + * When compute_input_grad=false (e.g. the input layer where x=states), + * the W^T @ dL_dz accumulation into dL_dx_dist is skipped entirely. + * This saves ~DIST_SIZE(STATE_DIM) registers and eliminates the + * corresponding shmem weight-row reads, reducing register pressure + * by ~51 registers at the last backward layer. + */ +__device__ void backward_fc_relu_tiled( + const float* dL_dh_dist, /* [DIST_SIZE(out_dim)] post-activation gradient */ + const float* h_dist, /* [DIST_SIZE(out_dim)] saved post-activation (for relu grad) */ + const float* x_dist, /* [DIST_SIZE(in_dim)] input activation */ + const float* __restrict__ W, /* [out_dim, in_dim] */ + float* grad_W, /* atomicAdd target */ + float* grad_b, /* atomicAdd target */ + float* dL_dx_dist, /* [DIST_SIZE(in_dim)] accumulated input gradient (NULL if !compute_input_grad) */ + int in_dim, int out_dim, + float* shmem_weights, + int lane_id, + bool compute_input_grad /* false for last layer (input = states, not trainable) */ +) { + /* First compute dL/dz = dL/dh * leaky_relu_derivative(h) + * and broadcast to all lanes via shared memory. */ + /* We need dL_dz[out_idx] as a scalar for each row, but it's distributed. + * For each output index i, dL_dz[i] = dL_dh_dist[i/32] * (h_dist[i/32] > 0 ? 1 : LEAKY_RELU_ALPHA) + * But only the lane that "owns" element i has the correct value. + * We'll compute it on-the-fly during the tile loop using warp_reduce_sum_all. */ + + for (int tile = 0; tile < (out_dim + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; tile++) { + int ts = tile * SHMEM_TILE_ROWS; + int tr = SHMEM_MIN(SHMEM_TILE_ROWS, out_dim - ts); + cooperative_load_tile(shmem_weights, W + ts * in_dim, tr * in_dim); + __syncwarp(0xFFFFFFFF); + + for (int r = 0; r < tr; r++) { + int out_idx = ts + r; + if (out_idx >= out_dim) break; + + /* Reconstruct dL_dz[out_idx] — only one lane has the right + * distributed element, so we extract via shuffle. */ + int owning_lane = out_idx % 32; + int dist_slot = out_idx / 32; + float my_dL_dh = dL_dh_dist[dist_slot]; + float my_h = h_dist[dist_slot]; + /* Get the owning lane's values */ + float dL_dh_val = __shfl_sync(0xFFFFFFFF, my_dL_dh, owning_lane); + float h_val = __shfl_sync(0xFFFFFFFF, my_h, owning_lane); + float relu_grad = (h_val > 0.0f) ? 1.0f : LEAKY_RELU_ALPHA; + float dz = dL_dh_val * relu_grad; + + const float* W_row = shmem_weights + r * in_dim; + + for (int j = lane_id; j < in_dim; j += 32) + atomicAdd(&grad_W[out_idx * in_dim + j], dz * x_dist[j / 32]); + if (lane_id == 0) + atomicAdd(&grad_b[out_idx], dz); + /* Skip W^T @ dL_dz when input gradient is not needed (last layer). + * Saves register pressure and eliminates shmem reads for dL_dx. */ + if (compute_input_grad) { + for (int j = lane_id; j < in_dim; j += 32) + dL_dx_dist[j / 32] += W_row[j] * dz; + } + } + __syncwarp(0xFFFFFFFF); + } +} + +/* ── Main backward kernel ──────────────────────────────────────────── */ + +extern "C" __global__ void dqn_backward_kernel( + /* ── Batch data ────────────────────────────────────────────── */ + const float* __restrict__ states, /* [B, STATE_DIM] */ + const int* __restrict__ actions, /* [B] factored action indices */ + const float* __restrict__ is_weights, /* [B] */ + + /* ── Saved activations from forward kernel ─────────────────── */ + const float* __restrict__ save_h_s1, /* [B, SHARED_H1] */ + const float* __restrict__ save_h_s2, /* [B, SHARED_H2] */ + const float* __restrict__ save_h_v, /* [B, VALUE_H] */ + const float* __restrict__ save_h_b0, /* [B, ADV_H] */ + const float* __restrict__ save_h_b1, /* [B, ADV_H] */ + const float* __restrict__ save_h_b2, /* [B, ADV_H] */ + const float* __restrict__ save_current_lp, /* [B, 3, NUM_ATOMS] */ + const float* __restrict__ save_projected, /* [B, 3, NUM_ATOMS] */ + + /* ── Online network weights (for W^T in backward) ──────────── */ + const float* __restrict__ w_s1, const float* __restrict__ b_s1_unused, + const float* __restrict__ w_s2, const float* __restrict__ b_s2_unused, + const float* __restrict__ w_v1, const float* __restrict__ b_v1_unused, + const float* __restrict__ w_v2, const float* __restrict__ b_v2_unused, + const float* __restrict__ w_b0fc, const float* __restrict__ b_b0fc_unused, + const float* __restrict__ w_b0out, const float* __restrict__ b_b0out_unused, + const float* __restrict__ w_b1fc, const float* __restrict__ b_b1fc_unused, + const float* __restrict__ w_b1out, const float* __restrict__ b_b1out_unused, + const float* __restrict__ w_b2fc, const float* __restrict__ b_b2fc_unused, + const float* __restrict__ w_b2out, const float* __restrict__ b_b2out_unused, + + /* ── Gradient output buffer (atomicAdd accumulation) ────────── */ + float* __restrict__ grad_buf, /* [TOTAL_PARAMS] */ + + /* ── Config ────────────────────────────────────────────────── */ + int batch_size +) { + extern __shared__ float shmem[]; + float* shmem_weights = shmem; + /* dL_dvalue accumulator in shared memory — placed after the weight tile. + * The backward kernel only uses shmem_weights from the weight tile region + * (SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM floats). The remaining shared memory + * (bias_tile, scratch, state_cache, support_cache) is unused. We repurpose + * NUM_ATOMS floats (51 × 4 = 204 bytes) for accumulating dL/dvalue_logits + * across the 3 branches, avoiding a 51-register array that would otherwise + * stay live across the entire branch backward loop (STEP 2). */ + float* shmem_dL_dvalue = shmem + SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM; + + int sample_id = blockIdx.x; + int lane_id = threadIdx.x; + if (sample_id >= batch_size) return; + + /* ── Load per-sample data ──────────────────────────────────── */ + int factored_action = actions[sample_id]; + int branch_action[3]; + branch_action[0] = factored_action / 9; + branch_action[1] = (factored_action % 9) / 3; + branch_action[2] = factored_action % 3; + float is_weight = is_weights[sample_id]; + float scale = is_weight / (float)NUM_BRANCHES; + + /* ── Load saved activations into distributed registers ─────── */ + float h_s1_dist[DIST_SIZE(SHARED_H1)]; + for (int i = 0; i < DIST_SIZE(SHARED_H1); i++) h_s1_dist[i] = 0.0f; + for (int i = lane_id; i < SHARED_H1; i += 32) + h_s1_dist[i / 32] = save_h_s1[sample_id * SHARED_H1 + i]; + + float h_s2_dist[DIST_SIZE(SHARED_H2)]; + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) h_s2_dist[i] = 0.0f; + for (int i = lane_id; i < SHARED_H2; i += 32) + h_s2_dist[i / 32] = save_h_s2[sample_id * SHARED_H2 + i]; + + float h_v_dist[DIST_SIZE(VALUE_H)]; + for (int i = 0; i < DIST_SIZE(VALUE_H); i++) h_v_dist[i] = 0.0f; + for (int i = lane_id; i < VALUE_H; i += 32) + h_v_dist[i / 32] = save_h_v[sample_id * VALUE_H + i]; + + float state_dist[DIST_SIZE(STATE_DIM)]; + for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) state_dist[i] = 0.0f; + for (int i = lane_id; i < STATE_DIM; i += 32) + state_dist[i / 32] = states[sample_id * STATE_DIM + i]; + + /* ════════════════════════════════════════════════════════════ + * STEP 1: Compute dL/dlogit for each branch from saved data + * + * For taken action a_d in branch d: + * softmax[k] = exp(current_lp[k]) + * dL/dlogit_d[k] = -projected[k] + softmax[k] + * + * dL/dvalue_logits is accumulated in shared memory (shmem_dL_dvalue) + * instead of a 51-element register array. This reduces peak register + * pressure during STEP 2 by ~51 registers, since dL_dlogit[NUM_ATOMS] + * is also live inside each branch iteration. + * ════════════════════════════════════════════════════════════ */ + + /* Zero-initialize shmem dL_dvalue accumulator (warp-cooperative) */ + for (int k = lane_id; k < NUM_ATOMS; k += 32) + shmem_dL_dvalue[k] = 0.0f; + __syncwarp(0xFFFFFFFF); + + /* Per-branch dL_dlogit — process and accumulate, then backward through layers */ + const int branch_sizes[3] = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE }; + const int branch_atom_counts[3] = { BRANCH_0_ATOMS, BRANCH_1_ATOMS, BRANCH_2_ATOMS }; + + /* Weight pointers for branches */ + const float* branch_fc_w[3] = { w_b0fc, w_b1fc, w_b2fc }; + const float* branch_out_w[3] = { w_b0out, w_b1out, w_b2out }; + + /* Gradient buffer offsets for branch FC and output layers */ + const int goff_bfc_w[3] = { GOFF_W_B0FC, GOFF_W_B1FC, GOFF_W_B2FC }; + const int goff_bfc_b[3] = { GOFF_B_B0FC, GOFF_B_B1FC, GOFF_B_B2FC }; + const int goff_bout_w[3] = { GOFF_W_B0OUT, GOFF_W_B1OUT, GOFF_W_B2OUT }; + const int goff_bout_b[3] = { GOFF_B_B0OUT, GOFF_B_B1OUT, GOFF_B_B2OUT }; + + /* Saved activation pointers for branch hidden */ + const float* branch_h_ptrs[3] = { + save_h_b0 + sample_id * ADV_H, + save_h_b1 + sample_id * ADV_H, + save_h_b2 + sample_id * ADV_H, + }; + + /* Accumulated gradient flowing back to h_s2 from all heads */ + float dL_dh_s2_dist[DIST_SIZE(SHARED_H2)]; + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) dL_dh_s2_dist[i] = 0.0f; + + /* ════════════════════════════════════════════════════════════ + * STEP 2: For each branch — backward through output + FC layers + * ════════════════════════════════════════════════════════════ */ + + for (int d = 0; d < NUM_BRANCHES; d++) { + int n_d = branch_sizes[d]; + int n_atoms = branch_atom_counts[d]; + int a_d = branch_action[d]; + + /* Load saved current_lp and projected for this branch */ + float dL_dlogit[NUM_ATOMS]; + for (int k = 0; k < NUM_ATOMS; k++) { + float lp = save_current_lp[sample_id * NUM_BRANCHES * NUM_ATOMS + d * NUM_ATOMS + k]; + float proj = save_projected[sample_id * NUM_BRANCHES * NUM_ATOMS + d * NUM_ATOMS + k]; + dL_dlogit[k] = -proj + expf(lp); + } + + /* Accumulate dL_dvalue in shared memory (reduces register pressure). + * Each lane adds its subset of atoms; no bank conflicts since + * adjacent lanes write adjacent addresses (stride-32 pattern). */ + for (int k = lane_id; k < NUM_ATOMS; k += 32) + shmem_dL_dvalue[k] += dL_dlogit[k]; + __syncwarp(0xFFFFFFFF); + + /* Load branch hidden activation into distributed registers */ + float h_bd_dist[DIST_SIZE(ADV_H)]; + for (int i = 0; i < DIST_SIZE(ADV_H); i++) h_bd_dist[i] = 0.0f; + for (int i = lane_id; i < ADV_H; i += 32) + h_bd_dist[i / 32] = branch_h_ptrs[d][i]; + + /* ── Backward through branch output layer ──────────────── */ + float dL_dh_bd_dist[DIST_SIZE(ADV_H)]; + for (int i = 0; i < DIST_SIZE(ADV_H); i++) dL_dh_bd_dist[i] = 0.0f; + + backward_output_tiled( + dL_dlogit, a_d, n_d, scale, + h_bd_dist, + branch_out_w[d], + grad_buf + goff_bout_w[d], + grad_buf + goff_bout_b[d], + dL_dh_bd_dist, + ADV_H, n_atoms, + shmem_weights, lane_id + ); + + /* ── Backward through branch FC (LeakyReLU + Linear) ──── */ + float dL_dh_s2_branch_dist[DIST_SIZE(SHARED_H2)]; + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) dL_dh_s2_branch_dist[i] = 0.0f; + + backward_fc_relu_tiled( + dL_dh_bd_dist, h_bd_dist, + h_s2_dist, + branch_fc_w[d], + grad_buf + goff_bfc_w[d], + grad_buf + goff_bfc_b[d], + dL_dh_s2_branch_dist, + SHARED_H2, ADV_H, + shmem_weights, lane_id, + true /* compute_input_grad: need dL/dh_s2 for shared layers */ + ); + + /* Accumulate into total dL/dh_s2 */ + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) + dL_dh_s2_dist[i] += dL_dh_s2_branch_dist[i]; + } + + /* ════════════════════════════════════════════════════════════ + * STEP 3: Backward through value output layer + * + * Load accumulated dL_dvalue from shared memory into registers. + * The shmem accumulator was populated during the branch loop + * (STEP 2), keeping 51 registers free during that loop. + * ════════════════════════════════════════════════════════════ */ + + /* Read back accumulated dL_dvalue from shared memory */ + float dL_dvalue[NUM_ATOMS]; + for (int k = 0; k < NUM_ATOMS; k++) + dL_dvalue[k] = shmem_dL_dvalue[k]; + + float dL_dh_v_dist[DIST_SIZE(VALUE_H)]; + for (int i = 0; i < DIST_SIZE(VALUE_H); i++) dL_dh_v_dist[i] = 0.0f; + + backward_value_out_tiled( + dL_dvalue, scale, + h_v_dist, + w_v2, + grad_buf + GOFF_W_V2, + grad_buf + GOFF_B_V2, + dL_dh_v_dist, + shmem_weights, lane_id + ); + + /* ════════════════════════════════════════════════════════════ + * STEP 4: Backward through value FC (LeakyReLU + Linear) + * ════════════════════════════════════════════════════════════ */ + + float dL_dh_s2_value_dist[DIST_SIZE(SHARED_H2)]; + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) dL_dh_s2_value_dist[i] = 0.0f; + + backward_fc_relu_tiled( + dL_dh_v_dist, h_v_dist, + h_s2_dist, + w_v1, + grad_buf + GOFF_W_V1, + grad_buf + GOFF_B_V1, + dL_dh_s2_value_dist, + SHARED_H2, VALUE_H, + shmem_weights, lane_id, + true /* compute_input_grad: need dL/dh_s2 for shared layers */ + ); + + /* Add value head contribution to dL/dh_s2 */ + for (int i = 0; i < DIST_SIZE(SHARED_H2); i++) + dL_dh_s2_dist[i] += dL_dh_s2_value_dist[i]; + + /* ════════════════════════════════════════════════════════════ + * STEP 5: Backward through shared_1 (LeakyReLU + Linear) + * ════════════════════════════════════════════════════════════ */ + + float dL_dh_s1_dist[DIST_SIZE(SHARED_H1)]; + for (int i = 0; i < DIST_SIZE(SHARED_H1); i++) dL_dh_s1_dist[i] = 0.0f; + + backward_fc_relu_tiled( + dL_dh_s2_dist, h_s2_dist, + h_s1_dist, + w_s2, + grad_buf + GOFF_W_S2, + grad_buf + GOFF_B_S2, + dL_dh_s1_dist, + SHARED_H1, SHARED_H2, + shmem_weights, lane_id, + true /* compute_input_grad: need dL/dh_s1 for input layer */ + ); + + /* ════════════════════════════════════════════════════════════ + * STEP 6: Backward through shared_0 (LeakyReLU + Linear) + * No dL/dx needed (input is states, not trainable). + * + * Pass compute_input_grad=false to skip W^T @ dL_dz accumulation. + * This eliminates the DIST_SIZE(STATE_DIM) register array that + * was previously allocated as dL_dstate_unused and immediately + * discarded, saving ~ceil(STATE_DIM/32) registers of pressure. + * ════════════════════════════════════════════════════════════ */ + + backward_fc_relu_tiled( + dL_dh_s1_dist, h_s1_dist, + state_dist, + w_s1, + grad_buf + GOFF_W_S1, + grad_buf + GOFF_B_S1, + (float*)NULL, /* no input gradient needed */ + STATE_DIM, SHARED_H1, + shmem_weights, lane_id, + false /* compute_input_grad: states are not trainable */ + ); +} + +/* ══════════════════════════════════════════════════════════════════════ + * GRADIENT NORM KERNEL (Phase 3a) + * + * Computes gradient L2 norm (sum of squares) via warp + block reduction + * and atomicAdd into a single output float. Must be followed by + * dqn_adam_update_kernel which reads the completed norm. + * + * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void dqn_grad_norm_kernel( + const float* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ + float* __restrict__ out_grad_norm, /* [1] output: sum of squares */ + int total_params +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + float g = (idx < total_params) ? grads[idx] : 0.0f; + float g2 = g * g; + + /* Warp-level reduction via shuffle (no shared memory) */ + for (int offset = 16; offset > 0; offset >>= 1) + g2 += __shfl_xor_sync(0xFFFFFFFF, g2, offset); + + /* Block-level cross-warp reduction via shared memory. + * Pad warp_sums to 2 elements per warp (stride=2) to avoid bank conflicts: + * warp 0 → bank 0, warp 1 → bank 2, ... warp 7 → bank 14. + * Without padding, sequential warp writes to warp_sums[0..7] map to + * banks 0..7 which is conflict-free for writes, but the __syncthreads() + * + read-back pattern benefits from padding to avoid false sharing + * between adjacent cache lines on H100's 128-byte L1 lines. */ + __shared__ float warp_sums[16]; /* 8 warps × 2 stride (padded) */ + int warp_id = threadIdx.x / 32; + int warp_lane = threadIdx.x % 32; + if (warp_lane == 0) warp_sums[warp_id * 2] = g2; + __syncthreads(); + + /* First warp reduces across warps using shuffles (pure register path) */ + if (warp_id == 0) { + float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + if (warp_lane == 0) + atomicAdd(out_grad_norm, val); + } +} + +/* ══════════════════════════════════════════════════════════════════════ + * ADAM UPDATE KERNEL (Phase 3b) + * + * Reads the COMPLETED gradient L2 norm from dqn_grad_norm_kernel, + * applies gradient clipping, then AdamW parameter update. + * Trivially parallel: 1 thread per parameter element. + * + * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void dqn_adam_update_kernel( + float* __restrict__ params, /* [TOTAL_PARAMS] flattened parameters */ + const float* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ + float* __restrict__ m, /* [TOTAL_PARAMS] first moment (Adam) */ + float* __restrict__ v, /* [TOTAL_PARAMS] second moment (Adam) */ + const float* __restrict__ grad_norm_sq, /* [1] completed sum of squares */ + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + float max_grad_norm, + const int* __restrict__ t_ptr, /* Adam step counter on device (for CUDA Graph) */ + int total_params +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_params) return; + + int t = *t_ptr; + float g = grads[idx]; + + /* Clip using the COMPLETED norm (no race — separate kernel launch) */ + float norm = sqrtf(*grad_norm_sq + 1e-12f); + float clip_scale = (norm > max_grad_norm) ? (max_grad_norm / norm) : 1.0f; + float clipped_g = g * clip_scale; + + /* Adam update */ + float beta1_t = 1.0f - powf(beta1, (float)t); + float beta2_t = 1.0f - powf(beta2, (float)t); + + float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g; + float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g * clipped_g; + m[idx] = m_i; + v[idx] = v_i; + + float m_hat = m_i / beta1_t; + float v_hat = v_i / beta2_t; + + /* AdamW weight decay (decoupled) */ + params[idx] -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * params[idx]); +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs new file mode 100644 index 000000000..cf1d228a4 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -0,0 +1,1595 @@ +#![allow(unsafe_code)] + +//! Fused CUDA DQN training — forward + loss + backward + Adam in 4 kernel launches. +//! +//! Replaces 2,100+ Candle kernel dispatches per training batch with 4 fused CUDA +//! kernels captured in a CUDA Graph for zero-overhead replay. +//! +//! ## Architecture +//! +//! All tensors are pre-allocated at construction (fixed shapes for CUDA Graph +//! compatibility). Weight pointers come from existing `DuelingWeightSet` + +//! `BranchingWeightSet` objects (no duplication). The only per-batch host work +//! is copying batch data into pre-allocated input buffers and reading back +//! the scalar loss + td_errors. +//! +//! ## CUDA Graph +//! +//! The training kernel sequence (zero_grad → forward+loss → backward → adam → +//! unflatten) is captured into a CUDA Graph on the first `train_step()` call. +//! Subsequent calls replay the graph with zero kernel-launch overhead. +//! +//! Only the Adam step counter (`t_buf`) and batch input data need updating +//! before each replay — both happen outside the captured graph. +//! +//! ## Kernel phases +//! +//! 1. **Forward + Loss** (`dqn_forward_loss_kernel`): 3 forward passes (online/states, +//! target/next_states, online/next_states for Double DQN) + C51 distributional +//! cross-entropy loss. Saves activations for backward pass. +//! 2. **Backward** (`dqn_backward_kernel`): Gradient through C51 → log_softmax → +//! dueling → linear layers. Accumulates into a flat gradient buffer via atomicAdd. +//! 3. **Grad norm** (`dqn_grad_norm_kernel`): Gradient L2 norm via warp + block +//! reduction + atomicAdd into a single output float. +//! 4. **Adam update** (`dqn_adam_update_kernel`): Reads completed norm, clips +//! gradients, applies AdamW update over flattened parameter view. +//! +//! ## Parameter layout +//! +//! The flat parameter/gradient/moment buffers use the same layout as the CUDA +//! `GOFF_*` defines: 20 weight tensors concatenated in order (w_s1, b_s1, w_s2, +//! b_s2, ..., w_b2out, b_b2out). See `compute_param_sizes()`. + +use std::sync::Arc; + +use candle_core::cuda_backend::cudarc; +use cudarc::driver::{ + CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg, +}; +use cudarc::nvrtc::Ptx; +use tracing::info; + +use super::gpu_weights::{BranchingWeightSetBf16, DuelingWeightSetBf16}; + +use crate::MLError; +use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; + +// ── CUDA Graph wrapper ────────────────────────────────────────────────────── + +/// Wrapper to send `CudaGraph` across thread boundaries. +/// +/// The graph is only launched on the same stream/context that created it. +/// The trainer is not shared across threads in practice. +struct SendSyncGraph(CudaGraph); + +// Safety: CudaGraph is bound to a specific CUDA context. The trainer +// that owns it is always used from the thread that created the context. +// No concurrent access occurs. +unsafe impl Send for SendSyncGraph {} +// Safety: same reasoning — single-owner, no concurrent launch calls. +unsafe impl Sync for SendSyncGraph {} + +// ── Configuration ─────────────────────────────────────────────────────────── + +/// Network dimensions and training hyperparameters for the fused CUDA trainer. +#[derive(Debug, Clone)] +pub struct GpuDqnTrainConfig { + /// Input state dimension (e.g., 48 or 56 with OFI). + pub state_dim: usize, + /// First shared hidden layer width (default: 256). + pub shared_h1: usize, + /// Second shared hidden layer width (default: 256). + pub shared_h2: usize, + /// Value head hidden width (default: 128). + pub value_h: usize, + /// Advantage/branch head hidden width (default: 128). + pub adv_h: usize, + /// Number of C51 distributional atoms (default: 51). + pub num_atoms: usize, + /// C51 minimum support value (default: -25.0). + pub v_min: f32, + /// C51 maximum support value (default: 25.0). + pub v_max: f32, + /// Branch 0 (exposure) action count (default: 5). + pub branch_0_size: usize, + /// Branch 1 (order) action count (default: 3). + pub branch_1_size: usize, + /// Branch 2 (urgency) action count (default: 3). + pub branch_2_size: usize, + /// Training batch size (fixed for CUDA Graph capture). + pub batch_size: usize, + /// Discount factor for Bellman projection. + pub gamma: f32, + /// Learning rate for Adam optimizer. + pub lr: f32, + /// Adam β1 (first moment decay). + pub beta1: f32, + /// Adam β2 (second moment decay). + pub beta2: f32, + /// Adam ε (numerical stability). + pub epsilon: f32, + /// Decoupled weight decay (AdamW). + pub weight_decay: f32, + /// Maximum gradient L2 norm for clipping. + pub max_grad_norm: f32, +} + +impl Default for GpuDqnTrainConfig { + fn default() -> Self { + Self { + state_dim: 48, + shared_h1: 256, + shared_h2: 256, + value_h: 128, + adv_h: 128, + num_atoms: 51, + v_min: -25.0, + v_max: 25.0, + branch_0_size: 5, + branch_1_size: 3, + branch_2_size: 3, + batch_size: 256, + gamma: 0.99, + lr: 3e-4, + beta1: 0.9, + beta2: 0.999, + epsilon: 1e-8, + weight_decay: 1e-5, + max_grad_norm: 10.0, + } + } +} + +/// Result of a single fused training step. +#[derive(Debug, Clone)] +pub struct FusedTrainResult { + /// Batch-mean weighted loss (for logging). + pub total_loss: f32, + /// Per-sample TD errors for PER priority update. + pub td_errors: Vec, + /// Pre-clip gradient L2 norm (for monitoring). + pub grad_norm: f32, +} + +// ── Parameter layout ──────────────────────────────────────────────────────── +// +// Matches GOFF_* / PARAM_* defines in dqn_training_kernel.cu exactly. +// 20 weight tensors in order: w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2, +// w_b0fc, b_b0fc, w_b0out, b_b0out, w_b1fc, b_b1fc, w_b1out, b_b1out, +// w_b2fc, b_b2fc, w_b2out, b_b2out. + +/// Compute the size (element count) of each of the 20 weight tensors. +fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; 20] { + [ + cfg.shared_h1 * cfg.state_dim, // w_s1 + cfg.shared_h1, // b_s1 + cfg.shared_h2 * cfg.shared_h1, // w_s2 + cfg.shared_h2, // b_s2 + cfg.value_h * cfg.shared_h2, // w_v1 + cfg.value_h, // b_v1 + cfg.num_atoms * cfg.value_h, // w_v2 + cfg.num_atoms, // b_v2 + cfg.adv_h * cfg.shared_h2, // w_b0fc + cfg.adv_h, // b_b0fc + cfg.branch_0_size * cfg.num_atoms * cfg.adv_h, // w_b0out + cfg.branch_0_size * cfg.num_atoms, // b_b0out + cfg.adv_h * cfg.shared_h2, // w_b1fc + cfg.adv_h, // b_b1fc + cfg.branch_1_size * cfg.num_atoms * cfg.adv_h, // w_b1out + cfg.branch_1_size * cfg.num_atoms, // b_b1out + cfg.adv_h * cfg.shared_h2, // w_b2fc + cfg.adv_h, // b_b2fc + cfg.branch_2_size * cfg.num_atoms * cfg.adv_h, // w_b2out + cfg.branch_2_size * cfg.num_atoms, // b_b2out + ] +} + +fn compute_total_params(cfg: &GpuDqnTrainConfig) -> usize { + compute_param_sizes(cfg).iter().sum() +} + +// ── Main struct ───────────────────────────────────────────────────────────── + +/// Fused CUDA DQN trainer — replaces Candle dispatch chain with 4 kernel launches. +/// +/// All GPU buffers are pre-allocated at construction. The trainer borrows +/// weight sets from the caller (no weight duplication). Only batch input +/// data is uploaded per step; outputs (loss, td_errors, grad_norm) are downloaded. +/// +/// On the first `train_step()`, the full kernel sequence is captured into a +/// CUDA Graph. Subsequent calls replay the graph for zero kernel-launch overhead. +#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug +pub struct GpuDqnTrainer { + config: GpuDqnTrainConfig, + stream: Arc, + + // ── Compiled kernels ──────────────────────────────────────────── + forward_loss_kernel: CudaFunction, + backward_kernel: CudaFunction, + grad_norm_kernel: CudaFunction, + adam_update_kernel: CudaFunction, + ema_kernel: CudaFunction, + f32_to_bf16_kernel: CudaFunction, + + // ── BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ── + // Allocated lazily on first train_step when weight sets are available. + // Synced after each Adam step (online) and after each EMA update (target). + online_dueling_bf16: Option, + online_branching_bf16: Option, + target_dueling_bf16: Option, + target_branching_bf16: Option, + + // ── Batch input buffers (uploaded per step) ───────────────────── + states_buf: CudaSlice, // [B, STATE_DIM] + next_states_buf: CudaSlice, // [B, STATE_DIM] + actions_buf: CudaSlice, // [B] + rewards_buf: CudaSlice, // [B] + dones_buf: CudaSlice, // [B] + is_weights_buf: CudaSlice, // [B] + + // ── Activation save buffers (forward → backward) ──────────────── + save_h_s1: CudaSlice, // [B, SHARED_H1] + save_h_s2: CudaSlice, // [B, SHARED_H2] + save_h_v: CudaSlice, // [B, VALUE_H] + save_h_b0: CudaSlice, // [B, ADV_H] + save_h_b1: CudaSlice, // [B, ADV_H] + save_h_b2: CudaSlice, // [B, ADV_H] + save_current_lp: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] + save_projected: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] + + // ── Forward output buffers ────────────────────────────────────── + per_sample_loss_buf: CudaSlice, // [B] + td_errors_buf: CudaSlice, // [B] + total_loss_buf: CudaSlice, // [1] + + // ── Backward / Adam buffers ───────────────────────────────────── + grad_buf: CudaSlice, // [TOTAL_PARAMS] gradient accumulator + params_buf: CudaSlice, // [TOTAL_PARAMS] flat online parameters + m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment + v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment + grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm + + // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─ + t_buf: CudaSlice, // [1] current Adam step + + // ── Training state ────────────────────────────────────────────── + adam_step: i32, + total_params: usize, + params_initialized: bool, + + // ── Shared memory size ────────────────────────────────────────── + shmem_bytes: usize, + + // ── CUDA Graph ────────────────────────────────────────────────── + training_graph: Option, + + // ── Consolidated transfer buffers ───────────────────────────── + /// Single staging buffer for batch upload consolidation. + /// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)] + /// One HtoD transfer replaces 6 separate PCIe round-trips. + upload_staging_buf: CudaSlice, + /// Total element count of the staging buffer. + upload_staging_len: usize, + + /// Single readback buffer for loss download consolidation. + /// Layout: [total_loss(1) | grad_norm(1) | td_errors(B)] + /// One DtoH transfer replaces 3 separate PCIe round-trips. + readback_buf: CudaSlice, + /// Pre-allocated host-side readback vec (avoids per-step allocation). + readback_host: Vec, + + /// Pre-allocated host-side staging vec for batch upload (avoids per-step malloc). + /// Reused via `.clear()` + `.extend_from_slice()` each step. + upload_staging_host: Vec, +} + +impl GpuDqnTrainer { + /// Configured batch size (fixed at construction for CUDA Graph compatibility). + pub fn batch_size(&self) -> usize { + self.config.batch_size + } + + /// Create a new fused DQN trainer with pre-allocated GPU buffers. + /// + /// Compiles all 4 kernels (forward+loss, backward, grad_norm, adam_update) + /// from a single NVRTC compilation. All buffers are allocated once (fixed + /// shapes for CUDA Graph compatibility). + pub fn new( + stream: Arc, + config: GpuDqnTrainConfig, + ) -> Result { + let b = config.batch_size; + let total_params = compute_total_params(&config); + + // ── Compile all 5 training kernels from same module ────────── + let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel) = + compile_training_kernels(&stream, &config)?; + + // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ── + let ema_kernel = compile_ema_kernel(&stream)?; + + // ── Allocate batch input buffers ──────────────────────────── + let states_buf = alloc_f32(&stream, b * config.state_dim, "states")?; + let next_states_buf = alloc_f32(&stream, b * config.state_dim, "next_states")?; + let actions_buf = alloc_i32(&stream, b, "actions")?; + let rewards_buf = alloc_f32(&stream, b, "rewards")?; + let dones_buf = alloc_f32(&stream, b, "dones")?; + let is_weights_buf = alloc_f32(&stream, b, "is_weights")?; + + // ── Allocate activation save buffers ──────────────────────── + let num_branches = 3; + let save_h_s1 = alloc_f32(&stream, b * config.shared_h1, "save_h_s1")?; + let save_h_s2 = alloc_f32(&stream, b * config.shared_h2, "save_h_s2")?; + let save_h_v = alloc_f32(&stream, b * config.value_h, "save_h_v")?; + let save_h_b0 = alloc_f32(&stream, b * config.adv_h, "save_h_b0")?; + let save_h_b1 = alloc_f32(&stream, b * config.adv_h, "save_h_b1")?; + let save_h_b2 = alloc_f32(&stream, b * config.adv_h, "save_h_b2")?; + let save_current_lp = alloc_f32( + &stream, + b * num_branches * config.num_atoms, + "save_current_lp", + )?; + let save_projected = alloc_f32( + &stream, + b * num_branches * config.num_atoms, + "save_projected", + )?; + + // ── Allocate forward output buffers ───────────────────────── + let per_sample_loss_buf = alloc_f32(&stream, b, "per_sample_loss")?; + let td_errors_buf = alloc_f32(&stream, b, "td_errors")?; + let total_loss_buf = alloc_f32(&stream, 1, "total_loss")?; + + // ── Allocate backward / Adam buffers ──────────────────────── + let grad_buf = alloc_f32(&stream, total_params, "grad_buf")?; + let params_buf = alloc_f32(&stream, total_params, "params_buf")?; + let m_buf = alloc_f32(&stream, total_params, "adam_m")?; + let v_buf = alloc_f32(&stream, total_params, "adam_v")?; + let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?; + let t_buf = alloc_i32(&stream, 1, "adam_t")?; + + // ── Allocate consolidated transfer buffers ───────────────── + // Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights + let upload_staging_len = b * config.state_dim * 2 + b * 4; // 2*B*SD + 4*B + let upload_staging_buf = alloc_f32(&stream, upload_staging_len, "upload_staging")?; + + // Readback: total_loss(1) + grad_norm(1) + td_errors(B) + let readback_len = 2 + b; + let readback_buf = alloc_f32(&stream, readback_len, "readback")?; + let readback_host = vec![0.0_f32; readback_len]; + + // Pre-allocate host staging vec for upload_batch (reused each step) + let upload_staging_host = vec![0.0_f32; upload_staging_len]; + + // ── Shared memory sizing ──────────────────────────────────── + let shmem_bytes = compute_shmem_bytes(&config); + + let batch_bytes = (b * config.state_dim * 2 + + b * 4 + + b * config.shared_h1 + + b * config.shared_h2 + + b * config.value_h + + b * config.adv_h * 3 + + b * num_branches * config.num_atoms * 2 + + b * 2 + 1) + * std::mem::size_of::(); + let optim_bytes = total_params * 4 * std::mem::size_of::(); // grad + params + m + v + + info!( + batch_size = b, + state_dim = config.state_dim, + total_params, + batch_alloc_mb = batch_bytes as f64 / (1024.0 * 1024.0), + optim_alloc_mb = optim_bytes as f64 / (1024.0 * 1024.0), + shmem_bytes, + "GpuDqnTrainer: all buffers allocated, 4 kernels compiled" + ); + + Ok(Self { + config, + stream, + forward_loss_kernel, + backward_kernel, + grad_norm_kernel, + adam_update_kernel, + ema_kernel, + f32_to_bf16_kernel, + online_dueling_bf16: None, + online_branching_bf16: None, + target_dueling_bf16: None, + target_branching_bf16: None, + states_buf, + next_states_buf, + actions_buf, + rewards_buf, + dones_buf, + is_weights_buf, + save_h_s1, + save_h_s2, + save_h_v, + save_h_b0, + save_h_b1, + save_h_b2, + save_current_lp, + save_projected, + per_sample_loss_buf, + td_errors_buf, + total_loss_buf, + grad_buf, + params_buf, + m_buf, + v_buf, + grad_norm_buf, + t_buf, + adam_step: 0, + total_params, + params_initialized: false, + shmem_bytes, + training_graph: None, + upload_staging_buf, + upload_staging_len, + readback_buf, + readback_host, + upload_staging_host, + }) + } + + // ═══════════════════════════════════════════════════════════════════ + // BF16 weight mirror management + // ═══════════════════════════════════════════════════════════════════ + + /// Ensure BF16 weight mirrors are allocated and synced from F32 originals. + /// + /// Called lazily on first `train_step()` or `forward_loss()`. Allocates + /// 4 BF16 mirror sets (online + target, dueling + branching) and syncs + /// initial F32 weights into them via the `f32_to_bf16_kernel`. + fn ensure_bf16_mirrors( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + if self.online_dueling_bf16.is_some() { + return Ok(()); + } + + // Allocate BF16 mirrors from F32 weight sets + let mut od_bf16 = DuelingWeightSetBf16::alloc_from(online_d, &self.stream)?; + let mut ob_bf16 = BranchingWeightSetBf16::alloc_from(online_b, &self.stream)?; + let mut td_bf16 = DuelingWeightSetBf16::alloc_from(target_d, &self.stream)?; + let mut tb_bf16 = BranchingWeightSetBf16::alloc_from(target_b, &self.stream)?; + + // Initial F32 → BF16 sync + od_bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?; + ob_bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?; + td_bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?; + tb_bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?; + + self.online_dueling_bf16 = Some(od_bf16); + self.online_branching_bf16 = Some(ob_bf16); + self.target_dueling_bf16 = Some(td_bf16); + self.target_branching_bf16 = Some(tb_bf16); + + info!("GpuDqnTrainer: BF16 weight mirrors allocated and synced (4 sets, 40 tensors)"); + Ok(()) + } + + /// Sync online BF16 mirrors from F32 weight tensors. + /// + /// Called inside the CUDA Graph after `unflatten_online_weights()` so that + /// the next graph replay's forward kernel reads updated BF16 weights. + /// 12 conversion kernel launches (one per dueling tensor) + 8 (branching). + fn sync_online_bf16( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + if let Some(ref mut bf16) = self.online_dueling_bf16 { + bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?; + } + if let Some(ref mut bf16) = self.online_branching_bf16 { + bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?; + } + Ok(()) + } + + /// Sync target BF16 mirrors from F32 weight tensors. + /// + /// Called after `target_ema_update()` so that the next forward kernel + /// reads updated BF16 target weights. + fn sync_target_bf16( + &mut self, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + if let Some(ref mut bf16) = self.target_dueling_bf16 { + bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?; + } + if let Some(ref mut bf16) = self.target_branching_bf16 { + bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?; + } + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // Full training step (CUDA Graph — capture once, replay many) + // ═══════════════════════════════════════════════════════════════════ + + /// Run a complete fused training step: forward+loss → backward → Adam. + /// + /// On the first call (or after `invalidate_training_graph()`), captures the + /// kernel sequence into a CUDA Graph. Subsequent calls replay the graph + /// with zero kernel-launch overhead. + /// + /// Per-step host work (OUTSIDE the graph): + /// - Upload batch data (host pointers change each step) + /// - Update Adam step counter `t_buf` + /// - Download results (loss, td_errors, grad_norm) + /// + /// Captured IN the graph (replayed via `graph.launch()`): + /// - Zero accumulators (memset_zeros) + /// - Forward+loss kernel + /// - Backward kernel + /// - Adam kernel + /// - Unflatten d2d copies (params_buf → individual weight tensors) + #[allow(clippy::too_many_arguments)] + pub fn train_step( + &mut self, + states: &[f32], + next_states: &[f32], + actions: &[i32], + rewards: &[f32], + dones: &[f32], + is_weights: &[f32], + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, + target_dueling: &DuelingWeightSet, + target_branching: &BranchingWeightSet, + ) -> Result { + let b = self.config.batch_size; + let sd = self.config.state_dim; + + // ── Validate ──────────────────────────────────────────────── + if states.len() != b * sd { + return Err(MLError::ModelError(format!( + "states length {} != batch_size({b}) × state_dim({sd})", + states.len() + ))); + } + if actions.len() != b { + return Err(MLError::ModelError(format!( + "actions length {} != batch_size({b})", + actions.len() + ))); + } + + // ── First call: flatten weights + allocate BF16 mirrors ────── + if !self.params_initialized { + self.flatten_online_weights(online_dueling, online_branching)?; + self.ensure_bf16_mirrors( + online_dueling, online_branching, + target_dueling, target_branching, + )?; + self.params_initialized = true; + } + + // ── Upload batch data (OUTSIDE graph — host pointers change) ─ + self.upload_batch(states, next_states, actions, rewards, dones, is_weights)?; + + // ── Update Adam step counter on device (OUTSIDE graph) ─────── + self.adam_step += 1; + self.stream + .memcpy_htod(&[self.adam_step], &mut self.t_buf) + .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?; + + // ── CUDA Graph: capture on first call, replay on subsequent ── + if let Some(ref graph) = self.training_graph { + graph.0.launch().map_err(|e| { + MLError::ModelError(format!("CUDA graph replay: {e}")) + })?; + } else { + self.capture_training_graph( + online_dueling, + online_branching, + )?; + } + + // ── Consolidate readback: gather 3 GPU buffers → 1 readback buf ── + // Layout: [total_loss(1) | grad_norm(1) | td_errors(B)] + // D2D copies are async on the same stream — no host sync until the + // single memcpy_dtoh at the end. + let readback_base = raw_device_ptr(&self.readback_buf, &self.stream); + let f32_bytes = std::mem::size_of::(); + + // total_loss → readback[0] + let loss_src = raw_device_ptr(&self.total_loss_buf, &self.stream); + dtod_copy(readback_base, loss_src, f32_bytes, &self.stream, 0, "readback_gather")?; + + // grad_norm → readback[1] + let norm_src = raw_device_ptr(&self.grad_norm_buf, &self.stream); + dtod_copy(readback_base + f32_bytes as u64, norm_src, f32_bytes, &self.stream, 1, "readback_gather")?; + + // td_errors → readback[2..2+B] + let td_src = raw_device_ptr(&self.td_errors_buf, &self.stream); + let td_bytes = b * f32_bytes; + dtod_copy(readback_base + (2 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 2, "readback_gather")?; + + // ── Single DtoH transfer ────────────────────────────────────── + self.stream + .memcpy_dtoh(&self.readback_buf, &mut self.readback_host) + .map_err(|e| MLError::ModelError(format!("DtoH readback: {e}")))?; + + // ── Unpack on CPU ───────────────────────────────────────────── + let total_loss = self.readback_host[0]; + let grad_norm_sq = self.readback_host[1]; + let td_errors = self.readback_host[2..2 + b].to_vec(); + + Ok(FusedTrainResult { + total_loss, + td_errors, + grad_norm: grad_norm_sq.sqrt(), + }) + } + + /// Run only the forward + C51 loss kernel (no backward/Adam). + /// + /// Useful for validation loss computation without updating weights. + /// This path is NOT graphed (called infrequently, different kernel set). + pub fn forward_loss( + &mut self, + states: &[f32], + next_states: &[f32], + actions: &[i32], + rewards: &[f32], + dones: &[f32], + is_weights: &[f32], + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, + target_dueling: &DuelingWeightSet, + target_branching: &BranchingWeightSet, + ) -> Result { + let b = self.config.batch_size; + let sd = self.config.state_dim; + + if states.len() != b * sd { + return Err(MLError::ModelError(format!( + "states length {} != batch_size({b}) × state_dim({sd})", + states.len() + ))); + } + if actions.len() != b { + return Err(MLError::ModelError(format!( + "actions length {} != batch_size({b})", + actions.len() + ))); + } + + // Ensure BF16 mirrors are ready (lazy init on first call) + self.ensure_bf16_mirrors( + online_dueling, online_branching, + target_dueling, target_branching, + )?; + + // Upload batch data + self.upload_batch(states, next_states, actions, rewards, dones, is_weights)?; + + // Zero total_loss accumulator + self.stream + .memset_zeros(&mut self.total_loss_buf) + .map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?; + + // Launch forward + loss kernel (reads BF16 weight mirrors) + self.launch_forward_loss()?; + + // ── Consolidate readback: gather 2 GPU buffers → 1 readback buf ── + // Layout: [total_loss(1) | td_errors(B)] — reuses the 2+B readback_buf + // D2D copies are async on the same stream — no host sync until the + // single memcpy_dtoh at the end. + let readback_base = raw_device_ptr(&self.readback_buf, &self.stream); + let f32_bytes = std::mem::size_of::(); + + // total_loss → readback[0] + let loss_src = raw_device_ptr(&self.total_loss_buf, &self.stream); + dtod_copy(readback_base, loss_src, f32_bytes, &self.stream, 0, "fwd_readback_gather")?; + + // td_errors → readback[1..1+B] + let td_src = raw_device_ptr(&self.td_errors_buf, &self.stream); + let td_bytes = b * f32_bytes; + dtod_copy(readback_base + f32_bytes as u64, td_src, td_bytes, &self.stream, 1, "fwd_readback_gather")?; + + // ── Single DtoH transfer ────────────────────────────────────── + self.stream + .memcpy_dtoh(&self.readback_buf, &mut self.readback_host) + .map_err(|e| MLError::ModelError(format!("DtoH fwd_readback: {e}")))?; + + // ── Unpack on CPU ───────────────────────────────────────────── + let total_loss = self.readback_host[0]; + let td_errors = self.readback_host[1..1 + b].to_vec(); + + Ok(FusedTrainResult { + total_loss, + td_errors, + grad_norm: 0.0, + }) + } + + // ═══════════════════════════════════════════════════════════════════ + // CUDA Graph capture and invalidation + // ═══════════════════════════════════════════════════════════════════ + + /// Capture the training kernel sequence into a CUDA Graph. + /// + /// Called on the first `train_step()` or after `invalidate_training_graph()`. + /// The captured graph includes: zero_grad → forward+loss → backward → adam → + /// unflatten (20 d2d copies). The graph.launch() executes the captured work. + fn capture_training_graph( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + // Begin stream capture — only work submitted from this thread on this + // stream is captured (THREAD_LOCAL mode, safe for single-stream use). + self.stream.begin_capture( + cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, + ).map_err(|e| MLError::ModelError(format!("CUDA graph begin_capture: {e}")))?; + + // Submit the training ops to the stream (captured into graph). + // MUST end capture even if submission fails. + let submit_result = + self.submit_training_ops(online_d, online_b); + + // End capture — instantiate the graph + let graph_result = self.stream.end_capture( + cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + + // Propagate submission error first + submit_result?; + + // Unwrap the graph — hard error if capture failed (no CPU fallback) + let graph = graph_result + .map_err(|e| MLError::ModelError(format!("CUDA graph end_capture: {e}")))? + .ok_or_else(|| MLError::ModelError( + "CUDA graph capture returned None — stream may not support capture".into() + ))?; + + // Launch the graph — this actually executes the captured work + graph.launch().map_err(|e| { + MLError::ModelError(format!("CUDA graph first launch: {e}")) + })?; + + info!( + "GpuDqnTrainer: CUDA graph captured and launched \ + (3 memsets + 4 kernels + 20 d2d unflatten)" + ); + self.training_graph = Some(SendSyncGraph(graph)); + Ok(()) + } + + /// Submit the capturable training kernel sequence to the stream. + /// + /// This is the inner loop extracted so it can be called both during + /// CUDA Graph capture and as a non-graphed fallback. + fn submit_training_ops( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + // ── Zero accumulators (capturable: memset_zeros uses cuMemsetD32Async) ─ + self.stream + .memset_zeros(&mut self.total_loss_buf) + .map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?; + self.stream + .memset_zeros(&mut self.grad_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_buf: {e}")))?; + self.stream + .memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("zero grad_norm: {e}")))?; + + // ── 1. Forward + Loss (reads BF16 weight mirrors) ───────── + self.launch_forward_loss()?; + + // ── 2. Backward (reads F32 weights for W^T) ───────────────── + self.launch_backward(online_d, online_b)?; + + // ── 3a. Gradient norm ─────────────────────────────────────── + self.launch_grad_norm()?; + + // ── 3b. Adam update ──────────────────────────────────────── + self.launch_adam_update()?; + + // ── 4. Unflatten: params_buf → individual F32 weight tensors ─ + // D2D copies are capturable (cuMemcpyDtoDAsync). Device pointers + // are stable — the graph replays writes to the same addresses. + self.unflatten_online_weights(online_d, online_b)?; + + // ── 5. Sync online BF16 mirrors from updated F32 tensors ───── + // 20 conversion kernel launches (capturable in CUDA Graph). + // Next graph replay's forward kernel reads updated BF16 weights. + self.sync_online_bf16(online_d, online_b)?; + + Ok(()) + } + + /// Discard the cached CUDA Graph and force re-flatten + re-capture. + /// + /// Call after: + /// - Target network EMA update (target weight pointers may change) + /// - Learning rate schedule change (lr is baked into the graph) + /// - Any external weight modification + /// + /// The next `train_step()` will re-capture a fresh graph. + pub fn invalidate_training_graph(&mut self) { + self.training_graph = None; + self.params_initialized = false; + // Drop BF16 mirrors — they'll be reallocated + synced on next train_step. + // This ensures the graph captures fresh BF16 pointers if weight sets change. + self.online_dueling_bf16 = None; + self.online_branching_bf16 = None; + self.target_dueling_bf16 = None; + self.target_branching_bf16 = None; + } + + /// Invalidate cached state after external weight modifications. + /// + /// Call after target network EMA update or any manual weight change. + /// Forces re-flatten on the next `train_step()`. + pub fn invalidate_params(&mut self) { + self.invalidate_training_graph(); + } + + // ═══════════════════════════════════════════════════════════════════ + // Batch upload helper + // ═══════════════════════════════════════════════════════════════════ + + /// Upload batch data to pre-allocated GPU buffers via consolidated staging. + /// + /// Packs all 6 arrays into a single contiguous host buffer, performs one + /// `memcpy_htod` to the GPU staging buffer, then scatters to individual + /// buffers via `memcpy_dtod_async`. This turns 6 PCIe round-trips into 1, + /// saving ~12-30us per batch on H100 PCIe Gen5. + /// + /// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)] + fn upload_batch( + &mut self, + states: &[f32], + next_states: &[f32], + actions: &[i32], + rewards: &[f32], + dones: &[f32], + is_weights: &[f32], + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let sd = self.config.state_dim; + + // ── Pack all data into pre-allocated host buffer (zero malloc) ── + self.upload_staging_host.clear(); + self.upload_staging_host.extend_from_slice(states); // B * SD f32s + self.upload_staging_host.extend_from_slice(next_states); // B * SD f32s + // Reinterpret i32 actions as f32 bits (same 4 bytes, no conversion) + for &a in actions { + self.upload_staging_host.push(f32::from_bits(a as u32)); + } + self.upload_staging_host.extend_from_slice(rewards); // B f32s + self.upload_staging_host.extend_from_slice(dones); // B f32s + self.upload_staging_host.extend_from_slice(is_weights); // B f32s + + // ── Single HtoD transfer ────────────────────────────────── + self.stream + .memcpy_htod(&self.upload_staging_host, &mut self.upload_staging_buf) + .map_err(|e| MLError::ModelError(format!("HtoD staging: {e}")))?; + + // ── Scatter from staging to individual buffers via DtoD ─── + let staging_base = raw_device_ptr(&self.upload_staging_buf, &self.stream); + let f32_size = std::mem::size_of::(); + let mut byte_offset: u64 = 0; + + // states: B * SD elements + let states_bytes = b * sd * f32_size; + let states_dst = raw_device_ptr(&self.states_buf, &self.stream); + dtod_copy(states_dst, staging_base + byte_offset, states_bytes, &self.stream, 0, "upload_scatter")?; + byte_offset += states_bytes as u64; + + // next_states: B * SD elements + let next_states_bytes = b * sd * f32_size; + let next_states_dst = raw_device_ptr(&self.next_states_buf, &self.stream); + dtod_copy(next_states_dst, staging_base + byte_offset, next_states_bytes, &self.stream, 1, "upload_scatter")?; + byte_offset += next_states_bytes as u64; + + // actions: B elements (reinterpreted as f32 in staging, copy raw bytes to i32 buf) + let actions_bytes = b * f32_size; // i32 and f32 are both 4 bytes + let actions_dst = raw_device_ptr_i32(&self.actions_buf, &self.stream); + dtod_copy(actions_dst, staging_base + byte_offset, actions_bytes, &self.stream, 2, "upload_scatter")?; + byte_offset += actions_bytes as u64; + + // rewards: B elements + let rewards_bytes = b * f32_size; + let rewards_dst = raw_device_ptr(&self.rewards_buf, &self.stream); + dtod_copy(rewards_dst, staging_base + byte_offset, rewards_bytes, &self.stream, 3, "upload_scatter")?; + byte_offset += rewards_bytes as u64; + + // dones: B elements + let dones_bytes = b * f32_size; + let dones_dst = raw_device_ptr(&self.dones_buf, &self.stream); + dtod_copy(dones_dst, staging_base + byte_offset, dones_bytes, &self.stream, 4, "upload_scatter")?; + byte_offset += dones_bytes as u64; + + // is_weights: B elements + let is_weights_bytes = b * f32_size; + let is_weights_dst = raw_device_ptr(&self.is_weights_buf, &self.stream); + dtod_copy(is_weights_dst, staging_base + byte_offset, is_weights_bytes, &self.stream, 5, "upload_scatter")?; + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // Kernel launch methods + // ═══════════════════════════════════════════════════════════════════ + + /// Launch the forward+loss kernel with BF16 weight mirrors. + /// + /// Argument order matches `dqn_forward_loss_kernel` in `dqn_training_kernel.cu`: + /// 6 batch data + 20 online BF16 weights + 20 target BF16 weights + 8 activation saves + /// + 3 outputs + 2 config = 59 args. + /// + /// BF16 mirrors must be populated before this call (via `ensure_bf16_mirrors`). + fn launch_forward_loss(&self) -> Result<(), MLError> { + let od = self.online_dueling_bf16.as_ref() + .ok_or_else(|| MLError::ModelError("BF16 online dueling mirrors not initialized".into()))?; + let ob = self.online_branching_bf16.as_ref() + .ok_or_else(|| MLError::ModelError("BF16 online branching mirrors not initialized".into()))?; + let td = self.target_dueling_bf16.as_ref() + .ok_or_else(|| MLError::ModelError("BF16 target dueling mirrors not initialized".into()))?; + let tb = self.target_branching_bf16.as_ref() + .ok_or_else(|| MLError::ModelError("BF16 target branching mirrors not initialized".into()))?; + + let b = self.config.batch_size; + let batch_size_i32 = b as i32; + let gamma = self.config.gamma; + + let launch_cfg = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: self.shmem_bytes as u32, + }; + + // Safety: argument order matches the extern "C" kernel signature exactly. + // All CudaSlice lifetimes are valid (owned by self). + // BF16 weight pointers (CudaSlice) match kernel's __nv_bfloat16* params. + // Grid/block dimensions match kernel expectations (1 warp per sample). + unsafe { + self.stream + .launch_builder(&self.forward_loss_kernel) + // ── Batch data (6) ────────────────────────────────── + .arg(&self.states_buf) + .arg(&self.next_states_buf) + .arg(&self.actions_buf) + .arg(&self.rewards_buf) + .arg(&self.dones_buf) + .arg(&self.is_weights_buf) + // ── Online network BF16 weights (20) ───────────────── + .arg(&od.w_s1) + .arg(&od.b_s1) + .arg(&od.w_s2) + .arg(&od.b_s2) + .arg(&od.w_v1) + .arg(&od.b_v1) + .arg(&od.w_v2) + .arg(&od.b_v2) + // Branch 0 (exposure) — DuelingWeightSet advantage slot + .arg(&od.w_a1) + .arg(&od.b_a1) + .arg(&od.w_a2) + .arg(&od.b_a2) + // Branch 1 (order) — BranchingWeightSet + .arg(&ob.w_bo1) + .arg(&ob.b_bo1) + .arg(&ob.w_bo2) + .arg(&ob.b_bo2) + // Branch 2 (urgency) — BranchingWeightSet + .arg(&ob.w_bu1) + .arg(&ob.b_bu1) + .arg(&ob.w_bu2) + .arg(&ob.b_bu2) + // ── Target network BF16 weights (20) ───────────────── + .arg(&td.w_s1) + .arg(&td.b_s1) + .arg(&td.w_s2) + .arg(&td.b_s2) + .arg(&td.w_v1) + .arg(&td.b_v1) + .arg(&td.w_v2) + .arg(&td.b_v2) + .arg(&td.w_a1) + .arg(&td.b_a1) + .arg(&td.w_a2) + .arg(&td.b_a2) + .arg(&tb.w_bo1) + .arg(&tb.b_bo1) + .arg(&tb.w_bo2) + .arg(&tb.b_bo2) + .arg(&tb.w_bu1) + .arg(&tb.b_bu1) + .arg(&tb.w_bu2) + .arg(&tb.b_bu2) + // ── Saved activations (8) ─────────────────────────── + .arg(&self.save_h_s1) + .arg(&self.save_h_s2) + .arg(&self.save_h_v) + .arg(&self.save_h_b0) + .arg(&self.save_h_b1) + .arg(&self.save_h_b2) + .arg(&self.save_current_lp) + .arg(&self.save_projected) + // ── Outputs (3) ───────────────────────────────────── + .arg(&self.per_sample_loss_buf) + .arg(&self.td_errors_buf) + .arg(&self.total_loss_buf) + // ── Config (2) ────────────────────────────────────── + .arg(&gamma) + .arg(&batch_size_i32) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_forward_loss_kernel launch: {e}")) + })?; + } + + Ok(()) + } + + /// Launch the backward kernel. + /// + /// Argument order matches `dqn_backward_kernel` in `dqn_training_kernel.cu`: + /// 3 batch data + 8 activation saves + 20 online weights + 1 grad_buf + 1 config = 33 args. + fn launch_backward( + &self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + let b = self.config.batch_size; + let batch_size_i32 = b as i32; + + let launch_cfg = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: self.shmem_bytes as u32, + }; + + // Safety: argument order matches the extern "C" kernel signature exactly. + // grad_buf was zeroed before this call. Weight pointers are read-only (for W^T). + unsafe { + self.stream + .launch_builder(&self.backward_kernel) + // ── Batch data (3) ────────────────────────────────── + .arg(&self.states_buf) + .arg(&self.actions_buf) + .arg(&self.is_weights_buf) + // ── Saved activations (8) ─────────────────────────── + .arg(&self.save_h_s1) + .arg(&self.save_h_s2) + .arg(&self.save_h_v) + .arg(&self.save_h_b0) + .arg(&self.save_h_b1) + .arg(&self.save_h_b2) + .arg(&self.save_current_lp) + .arg(&self.save_projected) + // ── Online weights for W^T (20, incl. unused biases) ─ + .arg(&online_d.w_s1) + .arg(&online_d.b_s1) // unused but positional + .arg(&online_d.w_s2) + .arg(&online_d.b_s2) + .arg(&online_d.w_v1) + .arg(&online_d.b_v1) + .arg(&online_d.w_v2) + .arg(&online_d.b_v2) + .arg(&online_d.w_a1) // branch 0 FC + .arg(&online_d.b_a1) + .arg(&online_d.w_a2) // branch 0 out + .arg(&online_d.b_a2) + .arg(&online_b.w_bo1) // branch 1 FC + .arg(&online_b.b_bo1) + .arg(&online_b.w_bo2) // branch 1 out + .arg(&online_b.b_bo2) + .arg(&online_b.w_bu1) // branch 2 FC + .arg(&online_b.b_bu1) + .arg(&online_b.w_bu2) // branch 2 out + .arg(&online_b.b_bu2) + // ── Gradient output (1) ───────────────────────────── + .arg(&self.grad_buf) + // ── Config (1) ────────────────────────────────────── + .arg(&batch_size_i32) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_backward_kernel launch: {e}")) + })?; + } + + Ok(()) + } + + /// Launch the gradient L2 norm reduction kernel. + /// + /// Argument order matches `dqn_grad_norm_kernel` in `dqn_training_kernel.cu`: + /// grads, out_grad_norm, total_params = 3 args. + /// + /// `grad_norm_buf` must be zeroed before this call (done in `submit_training_ops`). + fn launch_grad_norm(&self) -> Result<(), MLError> { + let tp = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, // uses static __shared__ warp_sums[8] + }; + + // Safety: argument order matches the extern "C" kernel signature exactly. + // grad_buf has size = total_params; grad_norm_buf has size = 1. + unsafe { + self.stream + .launch_builder(&self.grad_norm_kernel) + .arg(&self.grad_buf) + .arg(&self.grad_norm_buf) + .arg(&tp) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_grad_norm_kernel launch: {e}")) + })?; + } + + Ok(()) + } + + /// Launch the Adam update kernel with correct gradient clipping. + /// + /// Argument order matches `dqn_adam_update_kernel` in `dqn_training_kernel.cu`: + /// params, grads, m, v, grad_norm_sq, lr, beta1, beta2, epsilon, + /// weight_decay, max_grad_norm, t_ptr, total_params = 13 args. + /// + /// Must be launched AFTER `launch_grad_norm` so that `grad_norm_buf` + /// contains the completed sum of squares (no race condition). + /// `t_ptr` is a device buffer (not scalar) so the CUDA Graph can be + /// replayed with an updated step counter. + fn launch_adam_update(&self) -> Result<(), MLError> { + let tp = self.total_params as i32; + let blocks = ((self.total_params + 255) / 256) as u32; + + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + let lr = self.config.lr; + let beta1 = self.config.beta1; + let beta2 = self.config.beta2; + let epsilon = self.config.epsilon; + let weight_decay = self.config.weight_decay; + let max_grad_norm = self.config.max_grad_norm; + + // Safety: argument order matches the extern "C" kernel signature exactly. + // All buffers are pre-allocated with size = total_params. + // grad_norm_buf contains the completed norm from launch_grad_norm(). + // t_buf is a device pointer (kernel reads *t_ptr) for CUDA Graph compatibility. + unsafe { + self.stream + .launch_builder(&self.adam_update_kernel) + .arg(&self.params_buf) + .arg(&self.grad_buf) + .arg(&self.m_buf) + .arg(&self.v_buf) + .arg(&self.grad_norm_buf) + .arg(&lr) + .arg(&beta1) + .arg(&beta2) + .arg(&epsilon) + .arg(&weight_decay) + .arg(&max_grad_norm) + .arg(&self.t_buf) // device pointer — not baked scalar + .arg(&tp) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_adam_update_kernel launch: {e}")) + })?; + } + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // Weight flattening / unflattening (GPU d2d — zero host roundtrip) + // ═══════════════════════════════════════════════════════════════════ + + /// Copy 20 individual weight tensors into the flat `params_buf`. + /// + /// Order matches the GOFF_* layout in `dqn_training_kernel.cu`. + /// Pure device-to-device copies via `cuMemcpyDtoDAsync` — zero host roundtrip. + fn flatten_online_weights( + &self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + let sizes = compute_param_sizes(&self.config); + let dst_base = raw_device_ptr(&self.params_buf, &self.stream); + + // Ordered to match GOFF_* layout + let slices: [&CudaSlice; 20] = [ + &online_d.w_s1, &online_d.b_s1, + &online_d.w_s2, &online_d.b_s2, + &online_d.w_v1, &online_d.b_v1, + &online_d.w_v2, &online_d.b_v2, + &online_d.w_a1, &online_d.b_a1, // branch 0 = advantage slot + &online_d.w_a2, &online_d.b_a2, + &online_b.w_bo1, &online_b.b_bo1, // branch 1 + &online_b.w_bo2, &online_b.b_bo2, + &online_b.w_bu1, &online_b.b_bu1, // branch 2 + &online_b.w_bu2, &online_b.b_bu2, + ]; + + let mut byte_offset: u64 = 0; + for (i, slice) in slices.iter().enumerate() { + let num_bytes = sizes[i] * std::mem::size_of::(); + let src = raw_device_ptr(slice, &self.stream); + dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?; + byte_offset += num_bytes as u64; + } + + Ok(()) + } + + /// Copy flat `params_buf` back to 20 individual weight tensors. + /// + /// Called after each Adam step to sync the canonical weight storage. + /// Pure device-to-device copies via `cuMemcpyDtoDAsync` — zero host roundtrip. + /// + /// During CUDA Graph capture, these d2d copies are recorded and replayed. + /// The device pointers are stable (CudaSlice allocations don't move). + fn unflatten_online_weights( + &self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + ) -> Result<(), MLError> { + let sizes = compute_param_sizes(&self.config); + let src_base = raw_device_ptr(&self.params_buf, &self.stream); + + // Must match GOFF_* order exactly + let slices: [&CudaSlice; 20] = [ + &online_d.w_s1, &online_d.b_s1, + &online_d.w_s2, &online_d.b_s2, + &online_d.w_v1, &online_d.b_v1, + &online_d.w_v2, &online_d.b_v2, + &online_d.w_a1, &online_d.b_a1, + &online_d.w_a2, &online_d.b_a2, + &online_b.w_bo1, &online_b.b_bo1, + &online_b.w_bo2, &online_b.b_bo2, + &online_b.w_bu1, &online_b.b_bu1, + &online_b.w_bu2, &online_b.b_bu2, + ]; + + let mut byte_offset: u64 = 0; + for (i, slice) in slices.iter().enumerate() { + let num_bytes = sizes[i] * std::mem::size_of::(); + let dst = raw_device_ptr(slice, &self.stream); + dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?; + byte_offset += num_bytes as u64; + } + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════ + // GPU-native Polyak EMA target update + // ═══════════════════════════════════════════════════════════════════ + + /// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` + /// + /// Updates all 20 target weight tensors in-place from the corresponding + /// online weight tensors using the EMA kernel. Runs OUTSIDE the captured + /// CUDA Graph — device pointers are stable so the graph stays valid. + /// + /// Eliminates the reverse-sync → CPU Polyak → forward-sync round-trip + /// that previously required 120 D2D copies + 120 Candle ops per step. + pub fn target_ema_update( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, + tau: f32, + ) -> Result<(), MLError> { + let sizes = compute_param_sizes(&self.config); + + // Paired (target, online) slices in GOFF_* order (20 pairs) + let pairs: [(&CudaSlice, &CudaSlice); 20] = [ + (&target_d.w_s1, &online_d.w_s1), + (&target_d.b_s1, &online_d.b_s1), + (&target_d.w_s2, &online_d.w_s2), + (&target_d.b_s2, &online_d.b_s2), + (&target_d.w_v1, &online_d.w_v1), + (&target_d.b_v1, &online_d.b_v1), + (&target_d.w_v2, &online_d.w_v2), + (&target_d.b_v2, &online_d.b_v2), + (&target_d.w_a1, &online_d.w_a1), // branch 0 (exposure) + (&target_d.b_a1, &online_d.b_a1), + (&target_d.w_a2, &online_d.w_a2), + (&target_d.b_a2, &online_d.b_a2), + (&target_b.w_bo1, &online_b.w_bo1), // branch 1 (order) + (&target_b.b_bo1, &online_b.b_bo1), + (&target_b.w_bo2, &online_b.w_bo2), + (&target_b.b_bo2, &online_b.b_bo2), + (&target_b.w_bu1, &online_b.w_bu1), // branch 2 (urgency) + (&target_b.b_bu1, &online_b.b_bu1), + (&target_b.w_bu2, &online_b.w_bu2), + (&target_b.b_bu2, &online_b.b_bu2), + ]; + + for (i, (target_slice, online_slice)) in pairs.iter().enumerate() { + let n = sizes[i] as i32; + let blocks = ((sizes[i] + 255) / 256) as u32; + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: argument order matches the extern "C" dqn_ema_kernel signature. + // target and online CudaSlice buffers have size >= sizes[i]. + // Both belong to the same CUDA context as self.stream. + unsafe { + self.stream + .launch_builder(&self.ema_kernel) + .arg(*target_slice) + .arg(*online_slice) + .arg(&tau) + .arg(&n) + .launch(launch_cfg) + .map_err(|e| { + MLError::ModelError(format!("dqn_ema_kernel launch[{i}]: {e}")) + })?; + } + } + + // Sync target BF16 mirrors from updated F32 target weights + self.sync_target_bf16(target_d, target_b)?; + + Ok(()) + } +} + +// ── Compilation ───────────────────────────────────────────────────────────── + +/// Compile all 5 training kernels from a single NVRTC compilation. +/// +/// Returns (forward_loss, backward, grad_norm, adam_update, f32_to_bf16) kernel functions. +fn compile_training_kernels( + stream: &Arc, + config: &GpuDqnTrainConfig, +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { + let shmem_max_in_dim = config + .state_dim + .max(config.shared_h1) + .max(config.shared_h2); + + let shmem_tile_rows = compute_shmem_tile_rows(shmem_max_in_dim); + + // Inject all dimensions as compile-time constants via #define. + // These MUST precede common_device_functions.cuh (which has #error guards). + let dim_overrides = format!( + "#define STATE_DIM {state_dim}\n\ + #define MARKET_DIM 42\n\ + #define PORTFOLIO_DIM 3\n\ + #define SHARED_H1 {shared_h1}\n\ + #define SHARED_H2 {shared_h2}\n\ + #define VALUE_H {value_h}\n\ + #define ADV_H {adv_h}\n\ + #define NUM_ATOMS {num_atoms}\n\ + #define V_MIN ({v_min}f)\n\ + #define V_MAX ({v_max}f)\n\ + #define BRANCH_0_SIZE {b0}\n\ + #define BRANCH_1_SIZE {b1}\n\ + #define BRANCH_2_SIZE {b2}\n\ + #define BATCH_SIZE {batch}\n\ + #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n\ + #define SHMEM_TILE_ROWS {shmem_tile_rows}\n", + state_dim = config.state_dim, + shared_h1 = config.shared_h1, + shared_h2 = config.shared_h2, + value_h = config.value_h, + adv_h = config.adv_h, + num_atoms = config.num_atoms, + v_min = config.v_min, + v_max = config.v_max, + b0 = config.branch_0_size, + b1 = config.branch_1_size, + b2 = config.branch_2_size, + batch = config.batch_size, + ); + + let common_src = include_str!("common_device_functions.cuh"); + let kernel_src = include_str!("dqn_training_kernel.cu"); + let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}"); + + info!( + state_dim = config.state_dim, + shared_h1 = config.shared_h1, + shared_h2 = config.shared_h2, + value_h = config.value_h, + adv_h = config.adv_h, + num_atoms = config.num_atoms, + batch_size = config.batch_size, + shmem_tile_rows, + total_params = compute_total_params(config), + "GpuDqnTrainer: compiling 4 training kernels" + ); + + let context = stream.context(); + let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) + .map_err(|e| { + MLError::ModelError(format!("dqn_training_kernel compilation failed: {e}")) + })?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("dqn_training module load: {e}")) + })?; + + let forward_loss = module + .load_function("dqn_forward_loss_kernel") + .map_err(|e| { + MLError::ModelError(format!("dqn_forward_loss_kernel load: {e}")) + })?; + let backward = module + .load_function("dqn_backward_kernel") + .map_err(|e| { + MLError::ModelError(format!("dqn_backward_kernel load: {e}")) + })?; + let grad_norm = module + .load_function("dqn_grad_norm_kernel") + .map_err(|e| { + MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}")) + })?; + let adam_update = module + .load_function("dqn_adam_update_kernel") + .map_err(|e| { + MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")) + })?; + let f32_to_bf16 = module + .load_function("f32_to_bf16_kernel") + .map_err(|e| { + MLError::ModelError(format!("f32_to_bf16_kernel load: {e}")) + })?; + + info!("GpuDqnTrainer: 5 kernels compiled and loaded (incl. BF16 converter)"); + Ok((forward_loss, backward, grad_norm, adam_update, f32_to_bf16)) +} + +/// Compile the standalone Polyak EMA kernel. +/// +/// `ema_kernel(target, online, tau, n)`: `target[i] = (1-tau)*target[i] + tau*online[i]` +/// +/// This kernel is NOT captured in the CUDA Graph — it runs after graph replay +/// to blend online weights into the target network in-place. Device pointers +/// are stable across calls, so the training graph remains valid. +fn compile_ema_kernel(stream: &Arc) -> Result { + let src = r#" +extern "C" __global__ +void dqn_ema_kernel(float* __restrict__ target, + const float* __restrict__ online, + float tau, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + target[i] = (1.0f - tau) * target[i] + tau * online[i]; + } +} +"#; + let context = stream.context(); + let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) + .map_err(|e| MLError::ModelError(format!("dqn_ema_kernel compilation: {e}")))?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("dqn_ema module load: {e}")) + })?; + module.load_function("dqn_ema_kernel").map_err(|e| { + MLError::ModelError(format!("dqn_ema_kernel load: {e}")) + }) +} + +// ── Shared memory sizing ──────────────────────────────────────────────────── + +fn compute_shmem_tile_rows(shmem_max_in_dim: usize) -> usize { + let gpu_caps = crate::gpu::capabilities::cached_capabilities(); + let shmem_kb = crate::gpu::capabilities::max_shared_memory_kb(&gpu_caps.device_name); + let shmem_limit_bytes = (shmem_kb * 1024 * 80 / 100).max(49152); + let max_tile = shmem_limit_bytes / (4 * (shmem_max_in_dim + 1)); + let pow2 = (max_tile as u32).next_power_of_two() >> 1; + pow2.clamp(16, 256) as usize +} + +fn compute_shmem_bytes(config: &GpuDqnTrainConfig) -> usize { + let shmem_max_in_dim = config + .state_dim + .max(config.shared_h1) + .max(config.shared_h2); + let shmem_tile_rows = compute_shmem_tile_rows(shmem_max_in_dim); + // BF16 weight tiles are half the byte-width → 2× rows fit in the same region. + // Forward kernel bias tile is sized for the doubled BF16 tile rows. + let shmem_tile_rows_bf16 = 2 * shmem_tile_rows; + + let max_branch_size = config + .branch_0_size + .max(config.branch_1_size) + .max(config.branch_2_size); + + let weight_tile = shmem_tile_rows * shmem_max_in_dim; + let bias_tile = shmem_tile_rows_bf16; // sized for doubled BF16 tile rows + let scratch = max_branch_size * config.num_atoms + config.num_atoms; + // State vector cache: reloaded per-sample, used for all 3 forward passes + let state_cache = config.state_dim; + // C51 support atoms cache: z_j = V_MIN + j * DELTA_Z, constant per block + let support_cache = config.num_atoms; + + (weight_tile + bias_tile + scratch + state_cache + support_cache) + * std::mem::size_of::() +} + +// ── Device-to-device copy helpers ──────────────────────────────────────────── + +/// Extract raw CUDA device pointer (CUdeviceptr = u64) from a CudaSlice. +/// +/// Same pattern as `raw_device_ptr` in `gpu_weights.rs` — wraps the SyncOnDrop +/// guard in ManuallyDrop to skip read event recording (safe on same stream). +fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr(stream); + let _no_drop = std::mem::ManuallyDrop::new(guard); + ptr +} + +/// Extract raw CUDA device pointer from an i32 CudaSlice. +/// +/// Same pattern as `raw_device_ptr` but for `CudaSlice` (actions buffer). +fn raw_device_ptr_i32(slice: &CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr(stream); + let _no_drop = std::mem::ManuallyDrop::new(guard); + ptr +} + +/// Async device-to-device memcpy with error context. +fn dtod_copy( + dst: u64, + src: u64, + num_bytes: usize, + stream: &Arc, + idx: usize, + op: &str, +) -> Result<(), MLError> { + // Safety: caller guarantees src/dst are valid device pointers within + // the same CUDA context, num_bytes ≤ allocation size of both regions, + // and stream belongs to the same context. + unsafe { + cudarc::driver::result::memcpy_dtod_async(dst, src, num_bytes, stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("DtoD {op}[{idx}]: {e}")))?; + } + Ok(()) +} + +// ── Allocation helpers ────────────────────────────────────────────────────── + +fn alloc_f32( + stream: &Arc, + n: usize, + name: &str, +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} f32): {e}")) + }) +} + +fn alloc_i32( + stream: &Arc, + n: usize, + name: &str, +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} i32): {e}")) + }) +} diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 509f7714b..a32ec6a01 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -67,7 +67,7 @@ impl DQNAgentType { /// Batch greedy action selection — single forward pass per architecture head. /// /// Dispatches to the underlying DQN or RegimeConditionalDQN batch method. - pub fn batch_greedy_actions(&self, states: &Tensor) -> Result, MLError> { + pub fn batch_greedy_actions(&self, states: &Tensor) -> Result { match self { Self::Standard(agent) => agent.batch_greedy_actions(states), Self::RegimeConditional(agent) => agent.batch_greedy_actions(states), @@ -123,7 +123,7 @@ impl DQNAgentType { &self, states: &Tensor, temperature: f64, - ) -> Result, MLError> { + ) -> Result { match self { Self::Standard(agent) => agent.batch_softmax_actions(states, temperature), Self::RegimeConditional(agent) => agent.batch_softmax_actions(states, temperature), @@ -135,7 +135,7 @@ impl DQNAgentType { &self, states: &Tensor, temperature: f64, - ) -> Result, MLError> { + ) -> Result { match self { Self::Standard(agent) => { agent.batch_hierarchical_softmax_actions(states, temperature) @@ -592,6 +592,34 @@ impl DQNAgentType { self.memory().step(); } + /// Post-step bookkeeping for the fused CUDA training path. + /// + /// Delegates to `DQN::fused_post_step` -- increments training_steps, updates + /// PER priorities, steps beta annealing, and runs Polyak target update. + /// Only supported for Standard agents (not RegimeConditional). + pub fn fused_post_step(&mut self, td_errors: &[f32], indices: &[usize]) -> Result<(), crate::MLError> { + match self { + Self::Standard(agent) => agent.fused_post_step(td_errors, indices), + Self::RegimeConditional(_) => Err(crate::MLError::ModelError( + "Fused CUDA training not supported for RegimeConditional agent".into(), + )), + } + } + + /// Post-step bookkeeping WITHOUT target EMA (for GPU-native EMA path). + /// + /// Delegates to `DQN::fused_post_step_no_ema` -- increments training_steps, + /// updates PER priorities, steps beta annealing. Target EMA is done by the + /// GPU EMA kernel in `FusedTrainingCtx::run_full_step`. + pub fn fused_post_step_no_ema(&mut self, td_errors: &[f32], indices: &[usize]) -> Result<(), crate::MLError> { + match self { + Self::Standard(agent) => agent.fused_post_step_no_ema(td_errors, indices), + Self::RegimeConditional(_) => Err(crate::MLError::ModelError( + "Fused CUDA training not supported for RegimeConditional agent".into(), + )), + } + } + /// Flush GPU-accumulated max priority to CPU (single scalar readback per epoch). /// /// Call once at epoch boundary after all `update_priorities_gpu` calls in the diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs new file mode 100644 index 000000000..6ef0f56f8 --- /dev/null +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -0,0 +1,443 @@ +//! Fused CUDA Training Module +//! +//! High-performance H100-optimized training path that replaces 2,100+ Candle kernel +//! dispatches per batch with 3 fused CUDA kernels captured in a CUDA Graph: +//! +//! 1. **Forward + Loss kernel** -- shared trunk, 3 branching advantage heads, C51 distributional loss +//! 2. **Backward kernel** -- full backprop through all layers with gradient clipping +//! 3. **Adam optimizer kernel** -- fused parameter update with weight decay +//! +//! After the first batch, all 3 kernels are captured into a CUDA Graph and replayed +//! on subsequent steps with zero launch overhead. +//! +//! ## Weight Sync Architecture (GPU-native EMA) +//! +//! Device pointers in `CudaSlice` buffers are stable across in-place updates, so +//! the CUDA Graph remains valid across Polyak EMA target updates: +//! +//! 1. **After fused step**: online `CudaSlice` weights are already updated by the Adam kernel +//! 2. **GPU Polyak EMA**: `target[i] = (1-tau)*target[i] + tau*online[i]` via EMA kernel +//! 3. **No reverse/forward sync per step** -- VarMap sync deferred to epoch boundary +//! +//! This eliminates ~120 D2D copies + 120 Candle flatten/contiguous ops per step +//! (~13 GB wasted PCIe per epoch). + +use std::sync::Arc; + +use anyhow::Result; +use candle_core::{Device, Tensor}; +use tracing::info; + +use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; +use crate::cuda_pipeline::gpu_weights::{ + self, BranchingWeightSet, DuelingWeightSet, +}; +use crate::dqn::dqn::GpuTrainResult; +use crate::dqn::mixed_precision::training_dtype; +use crate::dqn::replay_buffer_type::BatchSample; +use super::config::DQNAgentType; +use super::DQNHyperparameters; + +/// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets. +/// +/// Weight sets are extracted from the Candle VarMap at initialization. Per-step +/// training operates entirely on GPU `CudaSlice` buffers: the Adam kernel updates +/// online weights, then the EMA kernel blends online into target weights in-place. +/// +/// VarMap sync is deferred to epoch boundary to avoid 120 D2D copies per step. +/// Device pointers are stable across in-place updates, so the CUDA Graph stays valid. +pub(crate) struct FusedTrainingCtx { + trainer: GpuDqnTrainer, + online_dueling: DuelingWeightSet, + online_branching: BranchingWeightSet, + target_dueling: DuelingWeightSet, + target_branching: BranchingWeightSet, + stream: Arc, + /// Batch size at creation time -- must match `current_batch_size` to reuse CUDA Graph. + batch_size: usize, + /// Steps since last VarMap sync (deferred to epoch boundary). + steps_since_varmap_sync: usize, +} + +impl FusedTrainingCtx { + /// Create a new fused training context. + /// + /// Extracts online + target weight sets from the DQN's Candle VarMaps into + /// `CudaSlice` buffers, compiles the 3 fused kernels + EMA kernel, and + /// allocates all pre-allocated buffers for CUDA Graph capture. + /// + /// Only valid when `device.is_cuda() && agent.is_using_branching()`. + pub(crate) fn new( + device: &Device, + agent: &DQNAgentType, + hyperparams: &DQNHyperparameters, + batch_size: usize, + ) -> Result { + let dqn = match agent { + DQNAgentType::Standard(d) => d, + DQNAgentType::RegimeConditional(_) => { + return Err(anyhow::anyhow!( + "Fused CUDA training requires Standard DQN agent (not RegimeConditional)" + )); + } + }; + + let branching_net = dqn.branching_q_network.as_ref().ok_or_else(|| { + anyhow::anyhow!("Fused CUDA training requires branching Q-network") + })?; + let branching_target = dqn.branching_target_network.as_ref().ok_or_else(|| { + anyhow::anyhow!("Fused CUDA training requires branching target network") + })?; + + // Get CudaStream from device + let cuda_dev = match device { + Device::Cuda(dev) => dev, + _ => return Err(anyhow::anyhow!("Fused CUDA training requires CUDA device")), + }; + let stream = cuda_dev.cuda_stream(); + + // Build config from DQN network dimensions + let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims(); + let config = GpuDqnTrainConfig { + state_dim: dqn.config.state_dim, + shared_h1, + shared_h2, + value_h, + adv_h, + num_atoms: dqn.config.num_atoms, + v_min: dqn.config.v_min, + v_max: dqn.config.v_max, + branch_0_size: dqn.config.num_actions, + branch_1_size: dqn.config.num_order_types, + branch_2_size: dqn.config.num_urgency_levels, + batch_size, + gamma: hyperparams.gamma as f32, + lr: hyperparams.learning_rate as f32, + beta1: 0.9, + beta2: 0.999, + epsilon: 1e-8, + weight_decay: 1e-5, + max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32, + }; + + // Extract weight sets from VarMaps (online + target) + let online_vars = branching_net.vars(); + let target_vars = branching_target.vars(); + + let online_dueling = + gpu_weights::extract_dueling_weights_branching(online_vars, &stream) + .map_err(|e| anyhow::anyhow!("Extract online dueling weights: {e}"))?; + let online_branching = + gpu_weights::extract_branching_weights(online_vars, &stream) + .map_err(|e| anyhow::anyhow!("Extract online branching weights: {e}"))?; + let target_dueling = + gpu_weights::extract_dueling_weights_branching(target_vars, &stream) + .map_err(|e| anyhow::anyhow!("Extract target dueling weights: {e}"))?; + let target_branching = + gpu_weights::extract_branching_weights(target_vars, &stream) + .map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?; + + // Create the fused trainer (compiles kernels, allocates buffers) + let trainer = GpuDqnTrainer::new(stream.clone(), config) + .map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?; + + info!( + batch_size, + "Fused CUDA training initialized: 4 kernels + EMA compiled, \ + ~291K params, CUDA Graph will capture on first step" + ); + + Ok(Self { + trainer, + online_dueling, + online_branching, + target_dueling, + target_branching, + stream, + batch_size, + steps_since_varmap_sync: 0, + }) + } + + /// Batch size this context was created for. + pub(crate) fn batch_size(&self) -> usize { + self.batch_size + } + + /// Steps since last VarMap sync. + pub(crate) fn steps_since_varmap_sync(&self) -> usize { + self.steps_since_varmap_sync + } + + /// Run one full fused training step. + /// + /// Executes the complete training cycle with GPU-native EMA: + /// 1. Extract flat arrays from `BatchSample` + /// 2. Forward + loss + backward + Adam via fused CUDA kernels (or CUDA Graph replay) + /// 3. GPU EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` (20 kernel launches) + /// 4. PER priority update + beta annealing + training_steps increment + /// 5. Return `GpuTrainResult` with GPU scalar tensors for monitoring + /// + /// No VarMap sync per step -- deferred to epoch boundary via `sync_to_varmap()`. + /// This eliminates ~120 D2D copies + 120 Candle ops per step (~13 GB/epoch saved). + pub(crate) fn run_full_step( + &mut self, + batch: &BatchSample, + agent: &mut DQNAgentType, + device: &Device, + ) -> Result { + let state_dim = agent.get_state_dim(); + let (states, next_states, actions, rewards, dones, is_weights) = + extract_batch_arrays(batch, state_dim); + + // Step 1: Fused forward + loss + backward + Adam + // The Adam kernel updates online CudaSlice weights in-place. + // Device pointers are stable -- CUDA Graph stays valid. + let fused_result = self.trainer.train_step( + &states, &next_states, &actions, &rewards, &dones, &is_weights, + &self.online_dueling, &self.online_branching, + &self.target_dueling, &self.target_branching, + ).map_err(|e| anyhow::anyhow!("Fused train_step: {e}"))?; + + // Step 2: GPU-native Polyak EMA target update + // Computes cosine-annealed tau and applies target[i] = (1-tau)*target[i] + tau*online[i] + // entirely on GPU. No VarMap round-trip. + { + let dqn = agent.as_standard_mut().ok_or_else(|| { + anyhow::anyhow!("Fused training requires Standard DQN agent") + })?; + + if dqn.config.use_soft_updates { + let tau = compute_cosine_annealed_tau( + dqn.get_training_steps(), + dqn.config.tau, + dqn.config.tau_final, + dqn.config.tau_anneal_steps, + ); + + self.trainer.target_ema_update( + &self.online_dueling, &self.online_branching, + &self.target_dueling, &self.target_branching, + tau as f32, + ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?; + } + } + + // Step 3: PER priority update + training_steps++ + beta annealing + // (no target EMA -- already done by GPU kernel above) + // td_errors from fused kernel are already f32 on CPU. + agent.fused_post_step_no_ema(&fused_result.td_errors, &batch.indices) + .map_err(|e| anyhow::anyhow!("Fused post_step_no_ema: {e}"))?; + + self.steps_since_varmap_sync += 1; + + // Step 4: Create GPU scalar tensors from fused results for monitoring code + Ok(GpuTrainResult { + loss_gpu: Tensor::new(fused_result.total_loss, device) + .map_err(|e| anyhow::anyhow!("Fused loss->Tensor: {e}"))?, + grad_norm_gpu: Tensor::new(fused_result.grad_norm, device) + .map_err(|e| anyhow::anyhow!("Fused grad_norm->Tensor: {e}"))?, + }) + } + + /// Sync CudaSlice weights back to VarMap (deferred -- called at epoch boundary). + /// + /// Copies: + /// - Online `CudaSlice -> VarMap` (reverse sync: Adam-updated weights) + /// - Target `CudaSlice -> VarMap` (reverse sync: EMA-blended weights) + /// + /// This keeps the Candle VarMap in sync for checkpointing, Q-value estimation + /// via `forward()`, and any non-fused code paths. + pub(crate) fn sync_to_varmap(&mut self, agent: &mut DQNAgentType) -> Result<()> { + let dqn = agent.as_standard_mut().ok_or_else(|| { + anyhow::anyhow!("VarMap sync requires Standard DQN agent") + })?; + + // Reverse sync online weights: CudaSlice -> VarMap + let online_vars = dqn.branching_q_network.as_ref() + .ok_or_else(|| anyhow::anyhow!("Missing branching Q-network"))? + .vars(); + gpu_weights::reverse_sync_dueling_weights_branching( + online_vars, &self.online_dueling, &self.stream, + ).map_err(|e| anyhow::anyhow!("Reverse sync online dueling: {e}"))?; + gpu_weights::reverse_sync_branching_weights( + online_vars, &self.online_branching, &self.stream, + ).map_err(|e| anyhow::anyhow!("Reverse sync online branching: {e}"))?; + + // Reverse sync target weights: CudaSlice -> VarMap + let target_vars = dqn.branching_target_network.as_ref() + .ok_or_else(|| anyhow::anyhow!("Missing branching target network"))? + .vars(); + gpu_weights::reverse_sync_dueling_weights_branching( + target_vars, &self.target_dueling, &self.stream, + ).map_err(|e| anyhow::anyhow!("Reverse sync target dueling: {e}"))?; + gpu_weights::reverse_sync_branching_weights( + target_vars, &self.target_branching, &self.stream, + ).map_err(|e| anyhow::anyhow!("Reverse sync target branching: {e}"))?; + + self.steps_since_varmap_sync = 0; + Ok(()) + } +} + +/// Cosine-annealed Polyak EMA coefficient (BYOL/MoCo v3 schedule). +/// +/// `tau(t) = tau_final - (tau_final - tau_base) * (cos(pi*t/T) + 1) / 2` +/// +/// - Early training: `tau ~ tau_base` (fast target adaptation) +/// - Late training: `tau ~ tau_final` (stability) +/// - `anneal_steps == 0`: fixed `tau_base` (no annealing) +fn compute_cosine_annealed_tau( + training_steps: u64, + tau_base: f64, + tau_final: f64, + anneal_steps: u64, +) -> f64 { + if anneal_steps > 0 { + let progress = (training_steps as f64 / anneal_steps as f64).min(1.0); + let cosine_factor = (std::f64::consts::PI * progress).cos(); + tau_final - (tau_final - tau_base) * (cosine_factor + 1.0) / 2.0 + } else { + tau_base + } +} + +/// Extract flat `f32`/`i32` arrays from a `BatchSample` for the fused CUDA trainer. +/// +/// Returns `(states, next_states, actions, rewards, dones, is_weights)`. +/// Rewards are converted from fixed-point (`i32 / 1_000_000`) to `f32`. +fn extract_batch_arrays( + batch: &BatchSample, + state_dim: usize, +) -> (Vec, Vec, Vec, Vec, Vec, Vec) { + let b = batch.experiences.len(); + let mut states = Vec::with_capacity(b * state_dim); + let mut next_states = Vec::with_capacity(b * state_dim); + let mut actions = Vec::with_capacity(b); + let mut rewards = Vec::with_capacity(b); + let mut dones = Vec::with_capacity(b); + + for exp in &batch.experiences { + states.extend_from_slice(&exp.state); + next_states.extend_from_slice(&exp.next_state); + actions.push(exp.action as i32); + rewards.push(exp.reward_f32()); + dones.push(if exp.done { 1.0_f32 } else { 0.0_f32 }); + } + + (states, next_states, actions, rewards, dones, batch.weights.clone()) +} + +/// GPU Q-value estimation -- called every 50 training steps for monitoring. +/// +/// Samples 10 experiences from the replay buffer, runs a forward pass through +/// the branching Q-network, and uses GPU-native reduction kernels for: +/// - Q-value divergence check (early stopping on runaway Q-values) +/// - Q-value statistics (min/max/mean/variance) +/// - Welford running mean accumulation (zero CPU sync) +/// +/// Shared by both `train_step_single_batch` and `train_step_with_accumulation`. +/// +/// Returns `Ok(avg_q)` or an error if GPU Q-value accumulation fails. +#[allow(clippy::indexing_slicing)] +pub(super) fn gpu_q_value_estimation( + agent: &mut DQNAgentType, + training_guard: &mut crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard, + device: &Device, +) -> Result { + use candle_core::IndexOp; + + let buffer = agent.memory(); + if buffer.len() == 0 { + return Err(anyhow::anyhow!( + "GPU Q-value estimation requires non-empty replay buffer" + )); + } + + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; + + let state_dim = agent.get_state_dim(); + let mut batch_tensor_opt: Option = None; + + // GPU PER path: use gpu_batch.states directly + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = Some( + gpu.states + .to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, + ); + } + + // CPU fallback: build tensor from experiences + if batch_tensor_opt.is_none() { + let mut state_data = Vec::with_capacity(sample_size * state_dim); + for exp in &batch_sample.experiences { + state_data.extend_from_slice(&exp.state); + } + if !state_data.is_empty() { + let tensor = Tensor::from_vec( + state_data, + (sample_size, state_dim), + device, + ) + .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? + .to_dtype(training_dtype(device)) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; + batch_tensor_opt = Some(tensor); + } + } + + let batch_tensor = batch_tensor_opt.ok_or_else(|| { + anyhow::anyhow!("GPU Q-value estimation: no tensor built (empty batch?)") + })?; + + // Suppress forward() monitoring to avoid to_vec2 GPU->CPU sync + agent.set_training_forward_active(true); + let batch_q_values = agent + .forward(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; + agent.set_training_forward_active(false); + + let num_actions = batch_q_values.dims().get(1).copied().unwrap_or(5); + + // Divergence check on first sample + let first_q = batch_q_values + .i(0) + .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; + let div_result = training_guard + .qvalue_divergence(&first_q, num_actions, 10000.0) + .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; + agent + .log_q_values_from_stats( + div_result.q_min, + div_result.q_max, + div_result.q_mean, + div_result.q_variance, + num_actions, + ) + .map_err(|e| { + tracing::info!("Early stopping (Q-value divergence): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Batch average via GPU reduction (one-step delay due to double-buffering) + let stats = training_guard + .qvalue_stats(&batch_q_values, sample_size, num_actions) + .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; + let cached_avg_q = stats.q_mean as f64; + + // Accumulate Q-value mean on GPU via Welford running mean (zero sync) + let avg_q_tensor = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; + training_guard + .accumulate_q_value(&avg_q_tensor) + .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; + + Ok(cached_avg_q) +} diff --git a/crates/ml/src/trainers/dqn/mod.rs b/crates/ml/src/trainers/dqn/mod.rs index bb58cc5b2..ba6dbbcd2 100644 --- a/crates/ml/src/trainers/dqn/mod.rs +++ b/crates/ml/src/trainers/dqn/mod.rs @@ -23,6 +23,7 @@ mod data_loading; mod early_stopping; pub(crate) mod financials; mod features; +mod fused_training; pub mod lr_scheduler; mod monitoring; mod risk; diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs deleted file mode 100644 index 0640974d9..000000000 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ /dev/null @@ -1,7122 +0,0 @@ -//! DQN Trainer Implementation -//! -//! Main training loop and execution logic for Deep Q-Network. - -use std::collections::VecDeque; -use std::path::PathBuf; -use std::sync::Arc; -use std::time::Duration; - -use anyhow::{Context, Result}; -use candle_core::{Device, Tensor}; -#[cfg(feature = "cuda")] -use candle_core::IndexOp; -use common::metrics::questdb_sink; -use common::metrics::training_metrics; -use common::CommonError; -use risk::drawdown_monitor::DrawdownMonitor; -use risk::safety::position_limiter::HybridPositionLimiter; -use risk::safety::PositionLimiterConfig; -use num_traits::ToPrimitive; -use crate::cuda_pipeline::DqnGpuData; -use rust_decimal::Decimal; -use tokio::sync::RwLock; -use tracing::{debug, info, warn}; -use uuid::Uuid; - -use crate::dqn::action_space::{ExposureLevel, FactoredAction}; -use crate::dqn::order_router::OrderRouter; -use ml_core::fill_simulator::{FillSimulator, FillResult}; -use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; -use crate::dqn::curiosity::CuriosityModule; -use crate::dqn::dqn::{DQN, DQNConfig}; -use crate::dqn::logging::{LoggingConfig, MetricsAggregator, log_epoch_start, log_epoch_end, log_training_progress}; -use crate::dqn::portfolio_tracker::PortfolioTracker; -use crate::dqn::regime_conditional::RegimeConditionalDQN; -use crate::dqn::reward::{RewardConfig, RewardFunction}; -use crate::dqn::target_update::convergence_half_life; -use crate::dqn::mixed_precision::training_dtype; -use crate::dqn::{Experience, TradingState}; -use crate::evaluation::metrics::calculate_var_cvar; -use crate::trainers::TargetUpdateMode; -use crate::TrainingMetrics; -use crate::features::microstructure_features::*; -use crate::labeling::triple_barrier::TripleBarrierEngine; -#[cfg(not(feature = "cuda"))] -use crate::labeling::triple_barrier::PricePoint; -#[cfg(not(feature = "cuda"))] -use crate::labeling::types::BarrierConfig; -use crate::memory_optimization::auto_batch_size::{AutoBatchSizer, BatchSizeConfig}; - -// Import from sibling modules -use super::config::{DQNAgentType, DQNHyperparameters}; -use super::financials::compute_epoch_financials; -use super::monitoring::TrainingMonitor; -use super::statistics::{FeatureStatistics, QValueStats}; -use super::EPISODE_LENGTH; -use crate::features::extraction::FeatureVector; - - -pub struct DQNTrainer { - /// DQN agent - agent: Arc>, - /// Training hyperparameters - pub(crate) hyperparams: DQNHyperparameters, - /// Device (GPU or CPU) - device: Device, - /// Training metrics - metrics: Arc>, - /// Loss history for plateau detection - loss_history: Vec, - /// Q-value history for floor detection - q_value_history: Vec, - /// Best validation loss achieved so far - best_val_loss: f64, - /// Validation data for computing validation loss - pub(crate) val_data: Vec<(FeatureVector, Vec)>, - /// Validation loss history for early stopping - val_loss_history: Vec, - /// C4: Sharpe history for Sharpe-based early stopping - sharpe_history: Vec, - /// C4: Best Sharpe ratio achieved so far (higher = better) - best_sharpe: f64, - /// Epoch with best validation loss - best_epoch: usize, - /// Step counter for gradient logging (logs every 10 steps) - gradient_logging_step: usize, - /// Original buffer_size before AutoReplaySizer (for gradient collapse warmup) - collapse_warmup_buffer_size: usize, - /// Portfolio state tracker for P&L-based rewards (Bug #2 fix) - pub portfolio_tracker: PortfolioTracker, - /// Feature normalization statistics (WAVE 3 FIX #2) - /// None during stats collection phase (epochs 0-10), Some during normalization phase (epochs 11+) - pub feature_stats: Option, - /// Sliding window of recent actions for reward calculation (max 100) - recent_actions: VecDeque, - /// Reward function for calculating rewards with recent actions - reward_fn: RewardFunction, - - // WAVE 16S: Adaptive Risk Management Components - /// Kelly criterion optimizer for position sizing (None if disabled) - pub(crate) kelly_optimizer: Option>, - /// Trade history for Kelly calculation (wins/losses) - pub(crate) trade_history: VecDeque, - /// Volatility tracker for epsilon adjustment (None if disabled) - pub(crate) volatility_returns: VecDeque, - /// PnL history for Sharpe calculation (max 1000 entries) - pub(crate) pnl_history: VecDeque, - - // Wave 16 Portfolio Features - /// Enable action masking (filters invalid actions before Q-value computation) - pub enable_action_masking: bool, - /// Maximum position size for action masking (default: 2.0) - pub max_position: f64, - /// Entropy regularizer for preventing policy collapse (None if disabled) - // entropy_regularizer removed — SAC-style entropy is computed directly on Q-value tensors in DQN::compute_loss_internal - /// Multi-asset portfolio tracker (None if single-asset mode) - pub multi_asset_portfolio: Option>, - /// Stress tester for robustness validation (None if disabled) - pub stress_tester: Option>, - - // Wave 16 Core Risk Features Integration - /// Drawdown monitor for tracking portfolio drawdowns (15% max drawdown) - pub drawdown_monitor: Option>, - /// Position limiter with 3-tier limits (±10.0 absolute, 1M notional, 10% concentration) - pub position_limiter: Option>, - /// Circuit breaker for stopping training on consecutive failures - pub circuit_breaker: Option>, - - // WAVE 3.10: Microstructure Feature Calculators (12 features) - pub(crate) micro_high_low_spread: HighLowSpread, - pub(crate) micro_vw_spread: VolumeWeightedSpread, - pub(crate) micro_tick_count: TickCount, - pub(crate) micro_inter_arrival: InterArrivalTime, - pub(crate) micro_buy_sell_imbalance: BuySellImbalance, - pub(crate) micro_kyle_lambda: KyleLambda, - pub(crate) micro_price_impact: PriceImpact, - pub(crate) micro_variance_ratio: VarianceRatio, - // Note: Roll Measure, Corwin-Schultz, Amihud, VPIN already exist in ml/src/microstructure/ - // We'll integrate those in the update logic - /// Track last timestamp for inter-arrival time calculation - last_timestamp_ns: u64, - /// Track last close price for microstructure calculations - pub(crate) last_close: f64, - - // WAVE 1.1: Triple Barrier Integration - /// Triple barrier engine for position exit labeling - triple_barrier: Arc>, - /// Active position tracker ID (None = no active position) - active_position_tracker: Option, - /// WAVE P3: Track previous simulated position for barrier tracking continuity - previous_simulated_position: f32, - - // WAVE 1.2: Safety Infrastructure Integration (8 Systems) - /// Loss history window for spike detection (size: 30) - safety_loss_history: VecDeque, - /// Loss plateau counter for anomaly detection - safety_loss_plateau_counter: usize, - /// Action counts for diversity monitoring (5 exposure actions) - safety_action_counts: std::collections::HashMap, - /// Memory manager for GPU OOM risk monitoring - safety_memory_manager: Arc>, - /// Safety enforcement level (Strict/Normal/Permissive) - safety_level: crate::safety::SafetyLevel, - /// Step counter for periodic safety checks - safety_step_counter: usize, - - /// Optional path to feature cache directory for faster hyperopt - pub(crate) feature_cache_dir: Option, - - /// Previous epoch's mean Q-value for overestimation detection - prev_epoch_q_mean: f64, - /// Current effective tau (may be temporarily increased if Q-values drift) - adaptive_tau: f64, - - /// WAVE 24 (Agent 17): Patience-based early stopping for anti-overfitting - early_stopping: super::early_stopping::EarlyStopping, - - /// WAVE 26 P0.6: Learning rate scheduler with warmup - lr_scheduler: super::lr_scheduler::LRScheduler, - - /// WAVE 26 P1.8: Curiosity module for intrinsic rewards (None if curiosity_weight = 0.0) - curiosity_module: Option, - - // WAVE 26 P1: Advanced DQN Features Integration - // P1.3: Sharpe Ratio Reward Component - /// Rolling buffer of returns for Sharpe ratio calculation (max: sharpe_window) - returns_history: VecDeque, - /// Sharpe reward weight (0.0 = disabled) - sharpe_weight: f64, - - // P1.6: Adaptive Dropout Scheduling - /// Optional dropout scheduler (None if disabled) - dropout_scheduler: Option, - - // P1.7: Hindsight Experience Replay (HER) - /// Optional HER buffer (None if her_ratio = 0.0) - her_buffer: Option>, - - // P1.9: Generalized Advantage Estimation (GAE) - /// Optional GAE calculator (None if disabled) - gae_calculator: Option, - - // P1.11: Noisy Network Sigma Scheduling - /// Optional noisy sigma scheduler (None if disabled) - noisy_sigma_scheduler: Option, - - // WAVE 30: Structured Logging Integration - /// Logging configuration for training metrics - logging_config: LoggingConfig, - /// Metrics aggregator for windowed training statistics - metrics_aggregator: MetricsAggregator, - - // WAVE 44: Multi-step returns integration - /// N-step buffer for multi-step TD learning (None if n_steps=1) - nstep_buffer: Option, - - /// Current effective batch size (may be reduced by OOM recovery) - current_batch_size: usize, - - /// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step) - cached_avg_q: f64, - /// Counter for Q-value estimation frequency (estimate every N training steps) - q_estimation_counter: u64, - - /// Pre-uploaded GPU training data (set once, reused across epochs) - gpu_data: Option, - - /// GPU portfolio simulator for CUDA-accelerated experience collection - #[cfg(feature = "cuda")] - gpu_portfolio_sim: Option, - - /// Raw cudarc targets buffer for CUDA kernel (parallel to candle Tensor in gpu_data) - #[cfg(feature = "cuda")] - targets_raw_cuda: Option>, - - /// Raw cudarc features buffer for CUDA experience kernel [num_bars * 42] - #[cfg(feature = "cuda")] - features_raw_cuda: Option>, - - /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b) - #[cfg(feature = "cuda")] - gpu_experience_collector: Option, - - /// GPU-fused epsilon-greedy action selector (eliminates argmax GPU->CPU sync barrier) - #[cfg(feature = "cuda")] - gpu_action_selector: Option, - - /// GPU training guard for zero-sync safety checks (loss clip, NaN, grad collapse) - #[cfg(feature = "cuda")] - training_guard: Option, - - /// GPU monitoring reducer — accumulates reward/action stats across kernel launches - #[cfg(feature = "cuda")] - gpu_monitoring: Option, - - /// Reusable GPU staging buffers for zero-alloc fold transitions - buffer_pool: Option, - - /// Double-buffered GPU data for zero-downtime fold transitions - double_buffer: Option, - - /// GPU-resident walk-forward data (entire dataset on GPU, per-fold views via index ranges) - #[cfg(feature = "cuda")] - gpu_walk_forward: Option, - - /// Multi-GPU configuration for data-parallel training (None = single GPU) - multi_gpu: Option, - - /// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs) - /// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork). - #[cfg(feature = "cuda")] - cached_n_episodes: Option, - - /// Pre-computed OFI features per bar (indexed by global bar position). - /// Populated during data loading when MBP-10 order book data is available. - /// Passed as `regime_features` in `TradingState::from_normalized()`. - /// Arc-shared to avoid 2.68 GB copy per hyperopt trial (41.9M × 8 × 8 bytes). - pub(crate) ofi_features: Option>, - /// Number of training bars (OFI offset for validation data). - /// val_data[i] corresponds to ofi_features[ofi_val_offset + i]. - pub(crate) ofi_val_offset: usize, - - // Phase C: Fill simulation and smart order routing - /// Fill simulator for order type-dependent execution modeling - fill_simulator: FillSimulator, - /// EMA of bar volatility (|close log return|) for OrderRouter routing decisions - vol_ema: f64, - /// Running median volatility estimate (slowly adapting EMA) - median_vol: f64, -} - -impl std::fmt::Debug for DQNTrainer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DQNTrainer") - .field("hyperparams", &self.hyperparams) - .finish_non_exhaustive() - } -} - -impl DQNTrainer { - /// Create new DQN trainer with hyperparameters and debug logging disabled - pub fn new(hyperparams: DQNHyperparameters) -> Result { - Self::new_with_debug(hyperparams, false) - } - - /// Create new DQN trainer with a specific compute device. - /// Used by hyperopt to share a single CUDA context across parallel trials. - pub fn new_with_device(hyperparams: DQNHyperparameters, device: Device) -> Result { - Self::new_internal(hyperparams, false, Some(device)) - } - - /// Create new DQN trainer with hyperparameters and configurable debug logging - /// - /// # Arguments - /// * `hyperparams` - DQN training hyperparameters - /// * `debug_logging` - Enable debug logging (REWARD_DEBUG, gradient norms, etc.) - pub fn new_with_debug(hyperparams: DQNHyperparameters, debug_logging: bool) -> Result { - Self::new_internal(hyperparams, debug_logging, None) - } - - fn new_internal(mut hyperparams: DQNHyperparameters, debug_logging: bool, override_device: Option) -> Result { - // Validate batch size is non-zero - if hyperparams.batch_size == 0 { - return Err(anyhow::anyhow!( - "Batch size must be greater than 0, got: {}", - hyperparams.batch_size - )); - } - - // WAVE 26 P2.2: Validate gradient_accumulation_steps > 0 - if hyperparams.gradient_accumulation_steps == 0 { - return Err(anyhow::anyhow!( - "gradient_accumulation_steps must be greater than 0, got: {}", - hyperparams.gradient_accumulation_steps - )); - } - - // Pre-compute hidden dims to get accurate model size for batch sizing. - // Align input_dim to 8 so the log matches the actual model dimensions. - // (device not yet created, so use the formula directly — CUDA always aligns) - let ofi_pre = hyperparams.mbp10_data_dir.is_some(); - let input_dim: usize = if ofi_pre { 56 } else { 48 }; // (53+7)&!7=56, (45+7)&!7=48 - let output_dim: usize = 5; - let hidden_dims: Vec = match hyperparams.hidden_dim_base { - Some(base) => { - let b = crate::cuda_pipeline::align_to_tensor_cores(base); - vec![b, b] // Constant-width: no tapering, no silently-discarded narrow layer - } - None => { - let caps = crate::gpu::capabilities::cached_capabilities(); - let base = crate::gpu::memory_profile::resolve_hidden_dim_base( - caps.free_vram_mb, - ); - let b = crate::cuda_pipeline::align_to_tensor_cores(base); - vec![b, b] // Constant-width: no tapering - } - }; - - // Compute accurate model size from actual network dimensions - let full_dims: Vec = std::iter::once(input_dim) - .chain(hidden_dims.iter().copied()) - .chain(std::iter::once(output_dim)) - .collect(); - let param_count = crate::gpu::memory_profile::network_param_count(&full_dims); - // FP32 params + AdamW state (2x for momentum/variance) + target network copy = ~4x - let model_overhead_mb = (param_count as f64 * 4.0 * 4.0) / (1024.0 * 1024.0); - info!( - "DQN network: {:?} → {} params, {:.1} MB overhead", - full_dims, param_count, model_overhead_mb - ); - - // Dynamic batch sizing: scale UP for larger GPUs, cap DOWN for smaller ones. - // Uses HardwareBudget for consistent sizing across DQN/PPO. - const STATIC_MAX_BATCH_SIZE: usize = 8192; - let max_safe_batch = match AutoBatchSizer::new() { - Ok(sizer) => { - let config = BatchSizeConfig { - model_memory_mb: model_overhead_mb, - safety_margin: 0.15, - ..BatchSizeConfig::default() - }; - let safe = sizer.max_safe_batch_size(&config); - info!( - "AutoBatchSizer: GPU VRAM ceiling = {} (configured: {})", - safe, hyperparams.batch_size - ); - safe - } - Err(e) => { - info!( - "AutoBatchSizer unavailable ({}), using static cap: {}", - e, STATIC_MAX_BATCH_SIZE - ); - STATIC_MAX_BATCH_SIZE - } - }; - - // Cap to VRAM ceiling from AutoBatchSizer (no separate scale-UP — - // AutoBatchSizer already accounts for model size and available VRAM) - if hyperparams.batch_size > max_safe_batch { - info!( - "DQN batch_size capped from {} → {} (VRAM ceiling)", - hyperparams.batch_size, max_safe_batch - ); - hyperparams.batch_size = max_safe_batch; - } - - // Use override device if provided (hyperopt shares one CUDA context), - // otherwise auto-detect GPU - let device = if let Some(dev) = override_device { - dev - } else { - Device::cuda_if_available(0) - .map_err(|e| anyhow::anyhow!("Failed to initialize device: {}", e))? - }; - - // Dynamic replay buffer sizing: scale replay capacity to available VRAM. - // Only activates when replay_buffer_vram_fraction > 0 and GPU is detected. - // Also computes the PER memory budget from actual VRAM — no hardcoded caps. - let original_buffer_size = hyperparams.buffer_size; // Save before AutoReplaySizer mutates it - let mut per_max_memory_bytes: usize = 4 * 1024 * 1024 * 1024; // CPU fallback: 4 GB - if hyperparams.replay_buffer_vram_fraction > 0.0 && device.is_cuda() { - use ml_core::memory_optimization::detect_gpu_hardware; - match detect_gpu_hardware() { - Ok(hw) => { - let raw_sd = if hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; - let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &device); - let replay_cfg = hw.optimal_replay_config( - aligned_sd, - hyperparams.replay_buffer_vram_fraction, - ); - per_max_memory_bytes = replay_cfg.per_max_buffer_bytes; - if replay_cfg.capacity != hyperparams.buffer_size { - info!( - "AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%, PER budget={:.0}MB)", - hyperparams.buffer_size, - replay_cfg.capacity, - hw.free_memory_mb, - hyperparams.replay_buffer_vram_fraction * 100.0, - per_max_memory_bytes as f64 / (1024.0 * 1024.0), - ); - hyperparams.buffer_size = replay_cfg.capacity; - } - } - Err(e) => { - info!( - "AutoReplaySizer unavailable ({}), using static buffer_size: {}", - e, hyperparams.buffer_size - ); - } - } - } else if device.is_cuda() { - // No auto-sizer, but still compute PER budget from VRAM - use ml_core::memory_optimization::detect_gpu_hardware; - if let Ok(hw) = detect_gpu_hardware() { - per_max_memory_bytes = hw.per_max_buffer_bytes(); - } else { - // GPU detection failed — keep CPU default (4 GB) - } - } else { - // CPU device — keep default 4 GB PER budget - } - - info!( - "Initializing DQN trainer on device: {:?}, using 5 exposure actions + OrderRouter", - if device.is_cuda() { "CUDA GPU" } else { "CPU" }, - ); - - // Auto-detect mixed precision capability based on GPU architecture - let mixed_precision_detected = if device.is_cuda() { - match crate::memory_optimization::auto_batch_size::detect_gpu_memory() { - Ok((_total, _free, ref name)) => { - let detected = crate::dqn::mixed_precision::detect_from_gpu_name(name); - match &detected { - Some(c) => info!("GPU mixed precision: {:?} enabled (GPU: {})", c.dtype, name), - None => info!("GPU mixed precision: disabled (GPU: {})", name), - } - detected - } - Err(_) => None, - } - } else { - None - }; - - // Create DQN configuration - // 42-feature architecture: OHLCV, technical, patterns, volume, time, statistical, regime - // Portfolio features (3) are populated via PortfolioTracker → 45 total state_dim - // With MBP-10 OFI features: +8 OFI features → 53 total - // - // Tensor core alignment: state_dim is rounded up to the next multiple of 8 - // (53→56, 45→48) so that cuBLAS dispatches BF16 HMMA instructions instead - // of falling back to scalar FMA. The extra columns are zero-padded at the - // data pipeline boundaries (GpuPreloadedData and train_batch CPU path). - let ofi_enabled = hyperparams.mbp10_data_dir.is_some(); - let raw_state_dim = if ofi_enabled { 53 } else { 45 }; - let state_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &device); - let config = DQNConfig { - state_dim, - num_actions: 5, // 5 exposure levels (Short100, Short50, Flat, Long50, Long100) - hidden_dims, - learning_rate: hyperparams.learning_rate, - gamma: hyperparams.gamma as f32, - epsilon_start: hyperparams.epsilon_start as f32, - epsilon_end: hyperparams.epsilon_end as f32, - epsilon_decay: hyperparams.epsilon_decay as f32, - replay_buffer_capacity: hyperparams.buffer_size, - collapse_warmup_capacity: original_buffer_size, - batch_size: hyperparams.batch_size, - min_replay_size: hyperparams.min_replay_size.min(hyperparams.buffer_size), // Cap at buffer_size to prevent deadlock - target_update_freq: hyperparams.target_update_frequency, // Use hyperparameter instead of hardcoded 1000 - use_double_dqn: true, - use_huber_loss: hyperparams.use_huber_loss, - huber_delta: hyperparams.huber_delta as f32, - leaky_relu_alpha: 0.01, // Standard LeakyReLU alpha (prevents dead neurons) - gradient_clip_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0), // Wave 11 Bug #1 fix: Dynamic clipping - - // WAVE 16 (Agent 36): Target update configuration - tau: hyperparams.tau, - tau_final: hyperparams.tau * 0.1, // Anneal to 10% of base tau - tau_anneal_steps: 100_000, - use_soft_updates: matches!(hyperparams.target_update_mode, TargetUpdateMode::Soft), - - // Rainbow DQN warmup period - warmup_steps: hyperparams.warmup_steps, - - // PER configuration - initial_capital: hyperparams.initial_capital as f64, - use_per: hyperparams.use_per, - use_gpu_replay_buffer: hyperparams.use_gpu_replay_buffer, - per_alpha: hyperparams.per_alpha, - per_beta_start: hyperparams.per_beta_start, - per_beta_max: 1.0, - per_beta_annealing_steps: hyperparams.epochs * 2000, // ~2000 steps/epoch (130k bars / ~64 batch_size) - per_max_memory_bytes, - - // Wave 2.1: Dueling Networks (ENABLED BY DEFAULT - Wave 6.4) - use_dueling: hyperparams.use_dueling, - dueling_hidden_dim: hyperparams.dueling_hidden_dim, - - // Wave 2.2: Multi-Step Returns (N-step TD) (ENABLED BY DEFAULT - Wave 6.4) - n_steps: hyperparams.n_steps, // Default: 3 (Rainbow DQN standard) - - // Wave 2.3: Distributional RL (C51) (ENABLED BY DEFAULT - Wave 6.4) - use_distributional: hyperparams.use_distributional, // Default: enabled (C51 distributional RL) - num_atoms: hyperparams.num_atoms, // Rainbow DQN standard: 51 atoms - v_min: hyperparams.v_min as f32, // Minimum value for distribution support - v_max: hyperparams.v_max as f32, // Maximum value for distribution support - - // Wave 2.4: Noisy Networks for Exploration (ENABLED BY DEFAULT - Wave 6.4) - use_noisy_nets: hyperparams.use_noisy_nets, // Default: enabled (replaces epsilon-greedy) - noisy_sigma_init: hyperparams.noisy_sigma_init, // Rainbow DQN standard: 0.5 - - // BUG #37 FIX: Q-value clipping (prevents step-level explosions) - enable_q_value_clipping: true, - q_value_clip_min: -500.0, - q_value_clip_max: 500.0, - - // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold (from hyperparams) - gradient_collapse_multiplier: hyperparams.gradient_collapse_multiplier, - gradient_collapse_patience: hyperparams.gradient_collapse_patience, - - use_cql: hyperparams.use_cql, - cql_alpha: hyperparams.cql_alpha, - use_iqn: hyperparams.use_qr_dqn, // Controlled by hyperopt - iqn_num_quantiles: hyperparams.num_quantiles, // Controlled by hyperopt - iqn_kappa: hyperparams.qr_kappa as f32, // Controlled by hyperopt (f64→f32) - iqn_embedding_dim: 64, // Fixed (not in search space) - use_branching: hyperparams.use_branching, - branch_hidden_dim: hyperparams.branch_hidden_dim, - use_regime_conditioning: true, // Always enable per-regime IS weights for branching loss - use_cvar_action_selection: false, - cvar_alpha: 0.05, - - #[allow(clippy::cast_possible_truncation)] - minimum_profit_factor: hyperparams.minimum_profit_factor as f32, - weight_decay: hyperparams.weight_decay, - dropout_rate: if hyperparams.enable_dropout_scheduler { hyperparams.dropout_initial } else { 0.0 }, - mixed_precision: hyperparams.mixed_precision.clone().or(mixed_precision_detected), - entropy_coefficient: hyperparams.entropy_coefficient.unwrap_or(0.01), - noisy_epsilon_floor: hyperparams.noisy_epsilon_floor.unwrap_or(0.0) as f32, // C2: NoisyNet handles exploration - use_count_bonus: hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, // C3 FIX: enable when coefficient > 0 - count_bonus_coefficient: hyperparams.count_bonus_coefficient.unwrap_or(0.0), - ..DQNConfig::default() - }; - - // Extract curiosity dims before config is moved into the agent - let curiosity_market_dim = config.curiosity_market_dim; - let curiosity_hidden_dim = config.curiosity_hidden_dim; - - // Create DQN agent - let agent = if hyperparams.enable_regime_qnetwork { - info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)"); - info!(" - Regime detection: ADX (raw index 40) + CUSUM direction (raw index 41)"); - info!(" - Classification: Trending (ADX>0.25), Volatile (ADX≤0.25 & |CUSUM|>0.7), Ranging (otherwise)"); - let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone()) - .map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?; - DQNAgentType::RegimeConditional(regime_agent) - } else { - info!("Creating standard DQN with single Q-network head"); - let standard_agent = DQN::new_on_device(config, device.clone()) - .map_err(|e| anyhow::anyhow!("Failed to create DQN agent: {}", e))?; - DQNAgentType::Standard(standard_agent) - }; - - - // Initialize portfolio tracker with $100k starting capital and 1 basis point spread - // Bug #2 fix: Portfolio features were hardcoded as [0.0, 0.0, 0.0] at line 1528 - let portfolio_tracker = PortfolioTracker::new( - hyperparams.initial_capital, // P2-A: Configurable capital - 0.0001, // 1 basis point spread (0.01%) - hyperparams.cash_reserve_percent, // Cash reserve requirement - ); - - // Initialize reward function with hyperparameter-driven configuration - // WAVE 10-A9 FIX: Wire hold_penalty_weight from hyperparameters to RewardConfig - // BUG #17 FIX: Add normalization and percentage-based P&L (enabled by default) - let reward_config = RewardConfig { - pnl_weight: Decimal::ONE, - risk_weight: Decimal::try_from(0.1).unwrap_or(Decimal::ZERO), - cost_weight: Decimal::ONE, // Bug #2 fix: 100% transaction cost weight (was 0.05, 20x too low) - hold_reward: Decimal::ZERO, // Flat position = no edge = zero reward (was +0.001, 20x trade PnL) - movement_threshold: Decimal::try_from(hyperparams.movement_threshold) - .unwrap_or(Decimal::ZERO), - hold_penalty_weight: Decimal::try_from(hyperparams.hold_penalty_weight) - .unwrap_or(Decimal::ZERO), // CRITICAL FIX - diversity_weight: Decimal::try_from(-0.1).unwrap_or(Decimal::ZERO), - enable_normalization: !hyperparams.use_dsr, // DSR replaces EMA normalizer - use_percentage_pnl: true, // Bug #17: Use percentage returns for scale-invariance - circuit_breaker_config: CircuitBreakerConfig::default(), - triple_barrier_profit_bonus: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO), - triple_barrier_stop_penalty: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO), - sharpe_weight: Decimal::ZERO, // WAVE 26 P1.3: Disabled by default - sharpe_window: 20, // WAVE 26 P1.3: Standard 20-period window - use_dsr: hyperparams.use_dsr, - dsr_eta: hyperparams.dsr_eta, - initial_capital: hyperparams.initial_capital as f64, - }; - let reward_fn = RewardFunction::new_with_debug(reward_config, debug_logging)?; - - // WAVE 1.1: Initialize triple barrier engine (max 1000 active trackers) - let triple_barrier = Arc::new(RwLock::new(TripleBarrierEngine::new(1000))); - info!("Triple barrier engine initialized with 1000 max trackers"); - - // WAVE 16S: Initialize Kelly optimizer if enabled - let kelly_optimizer = if hyperparams.enable_kelly_sizing { - use crate::risk::kelly_optimizer::{KellyCriterionOptimizer, KellyOptimizerConfig}; - let kelly_config = KellyOptimizerConfig { - max_fraction: hyperparams.kelly_max_fraction, - min_fraction: 0.01, - lookback_period: 252, - confidence_threshold: 0.6, - volatility_adjustment: true, - drawdown_protection: true, - }; - let optimizer = KellyCriterionOptimizer::new(kelly_config) - .map_err(|e| anyhow::anyhow!("Failed to create Kelly optimizer: {}", e))?; - info!("Kelly optimizer enabled (fractional={}, max={})", - hyperparams.kelly_fractional, hyperparams.kelly_max_fraction); - Some(Arc::new(optimizer)) - } else { - None - }; - - // Wave 16 Portfolio Features: Initialize action masking, entropy regularization, and stress testing - let enable_action_masking = hyperparams.enable_action_masking; - let max_position = hyperparams.max_position_absolute; // BLOCKER #2: Use hyperopt-tunable position limit - - // Entropy regularization: SAC-style computed directly on Q-values in DQN::compute_loss_internal - if hyperparams.enable_entropy_regularization { - let coeff = hyperparams.entropy_coefficient.unwrap_or(0.01); - info!("Entropy regularization enabled (coefficient={coeff:.4}, applied to Q-value softmax in loss)"); - } - - // Multi-asset portfolio tracking (disabled -- single-asset is the current production mode) - // - // When expanding to multi-asset trading: - // 1. Add `enable_multi_asset: bool` to DQNHyperparams (default false). - // 2. Initialize MultiAssetPortfolioTracker here when the flag is set. - // 3. Wire portfolio state into the DQN observation: expand state_dim to - // include per-asset position, PnL, and correlation features so the - // agent can learn cross-asset hedging and allocation. - let multi_asset_portfolio: Option> = None; - - // Stress testing for robustness validation - // Initialized as None here; call init_stress_tester() after construction - // to resolve the circular dependency (DQNStressTester needs a DQNTrainer). - let stress_tester: Option> = None; - - if enable_action_masking { - info!( - "Action masking enabled (max_position=±{:.1}, 30-50% filtering expected)", - max_position - ); - } else { - info!("Action masking disabled (all 5 exposure levels available)"); - } - - // Wave 16 Core Risk Features: Initialize drawdown monitor, position limiter, circuit breaker - // These are ALWAYS enabled by default for production safety - - // 1. Drawdown Monitor (15% max drawdown, alerts at 10%, 12.5%, 15%) - let drawdown_monitor = { - // DrawdownMonitor will be configured in first training step - // Config will be applied via async configure_alerts() in train_epoch - info!("Drawdown monitor enabled (thresholds: 10%, 12.5%, 15%)"); - Some(Arc::new(DrawdownMonitor::new())) - }; - - // 2. Position Limiter (3-tier limits: ±10.0 absolute, 1M notional, 10% concentration) - let position_limiter = { - let config = PositionLimiterConfig { - enabled: true, - cache_ttl: Duration::from_secs(60), - rpc_check_threshold_percent: 0.8, - max_position_per_symbol: 10.0, // ±10.0 absolute position limit - max_order_value: 1_000_000.0, // $1M notional limit - max_daily_loss: 0.10, // 10% concentration limit - }; - let limiter = HybridPositionLimiter::new(config); - info!("Position limiter enabled (abs=±10.0, notional=$1M, concentration=10%)"); - Some(Arc::new(limiter)) - }; - - // 3. Circuit Breaker (5 consecutive failures, 60s cooldown) - let circuit_breaker = { - let config = CircuitBreakerConfig { - failure_threshold: 5, - success_threshold: 3, - timeout_duration: Duration::from_secs(60), - half_open_max_calls: 2, - }; - let breaker = CircuitBreaker::new(config); - info!("Circuit breaker enabled (threshold=5 failures, cooldown=60s)"); - Some(Arc::new(breaker)) - }; - - // WAVE 24: Capture patience before hyperparams is moved - let early_stopping_patience = hyperparams.gradient_collapse_patience; - - // WAVE 26 P1: Initialize advanced DQN features - // P1.3: Sharpe ratio reward component - let sharpe_weight = hyperparams.sharpe_weight; - let sharpe_window = hyperparams.sharpe_window; - - // P1.6: Adaptive dropout scheduler - let dropout_scheduler = hyperparams.enable_dropout_scheduler.then(|| { - use crate::dqn::network::DropoutScheduler; - info!("Dropout scheduler enabled (initial={}, final={}, steps={})", - hyperparams.dropout_initial, hyperparams.dropout_final, hyperparams.dropout_anneal_steps); - DropoutScheduler::new( - hyperparams.dropout_initial, - hyperparams.dropout_final, - hyperparams.dropout_anneal_steps, - ) - }); - - // P1.7: Hindsight Experience Replay (HER) - let her_buffer = (hyperparams.her_ratio > 0.0) - .then(|| { - use crate::dqn::hindsight_replay::{HindsightReplayBuffer, HindsightReplayConfig, HindsightStrategy}; - use crate::dqn::prioritized_replay::PrioritizedReplayConfig; - let her_strategy = match hyperparams.her_strategy.as_str() { - "final" => HindsightStrategy::Final, - _ => HindsightStrategy::Future, // Default to Future - }; - let config = HindsightReplayConfig { - base_config: PrioritizedReplayConfig { - capacity: hyperparams.buffer_size, - ..Default::default() - }, - her_ratio: hyperparams.her_ratio, - her_strategy, - goal_dim: 1, // Single goal dimension for trading (target return) - k_future: 4, // Sample 4 future goals for Future strategy - batch_size: hyperparams.batch_size, - }; - info!("HER buffer enabled (ratio={}, strategy={:?}, capacity={})", - hyperparams.her_ratio, her_strategy, hyperparams.buffer_size); - HindsightReplayBuffer::new(config) - .map(Arc::new) - .map_err(|e| anyhow::anyhow!("Failed to create HER buffer: {}", e)) - }) - .transpose()?; - - // P1.9: Generalized Advantage Estimation (GAE) - let gae_calculator = hyperparams.enable_gae.then(|| { - use crate::dqn::gae::GAECalculator; - info!("GAE calculator enabled (lambda={}, gamma={})", - hyperparams.gae_lambda, hyperparams.gamma); - GAECalculator::new(hyperparams.gae_lambda, hyperparams.gamma) - }); - - // P1.11: Noisy network sigma scheduling - let noisy_sigma_scheduler = hyperparams.enable_noisy_sigma_scheduler.then(|| { - use crate::dqn::noisy_sigma_scheduler::NoisySigmaScheduler; - info!("Noisy sigma scheduler enabled (initial={}, final={}, steps={})", - hyperparams.noisy_sigma_initial, hyperparams.noisy_sigma_final, hyperparams.noisy_sigma_anneal_steps); - NoisySigmaScheduler::new( - hyperparams.noisy_sigma_initial, - hyperparams.noisy_sigma_final, - hyperparams.noisy_sigma_anneal_steps, - ) - }); - - // WAVE 26 P1.8: Initialize curiosity module if curiosity_weight > 0 - // Must be created BEFORE hyperparams and device are moved - let curiosity_module = (hyperparams.curiosity_weight > 0.0) - .then(|| { - CuriosityModule::new( - device.clone(), - 0.001, // Forward model learning rate - 2.0, // Max curiosity reward (clip to prevent noise exploitation) - curiosity_market_dim, - curiosity_hidden_dim, - 3, // Action categories (Short/Flat/Long) - ).map_err(|e| anyhow::anyhow!("Failed to create curiosity module: {}", e)) - }) - .transpose()?; - - // WAVE 26 P0.6: Initialize learning rate scheduler with warmup - // Must be created BEFORE hyperparams is moved - let lr_scheduler = { - use super::lr_scheduler::LRScheduler; - LRScheduler::new( - hyperparams.learning_rate, - hyperparams.warmup_steps, - hyperparams.lr_decay_type, - ) - }; - - // WAVE 44: Initialize n-step buffer if n_steps > 1 - let nstep_buffer = (hyperparams.n_steps > 1).then(|| { - info!("🎯 Multi-step returns ENABLED: n_steps={}, gamma={}", - hyperparams.n_steps, hyperparams.gamma); - crate::dqn::nstep_buffer::NStepBuffer::new( - hyperparams.n_steps, - hyperparams.gamma - ) - }); - - // Capture values before hyperparams is moved into Self - let initial_batch_size = hyperparams.batch_size; - let base_tau = hyperparams.tau; - - // GPU pipeline: pre-allocate staging buffers on CUDA devices - let buffer_pool = device.is_cuda().then(|| { - info!("GpuBufferPool: pre-allocated staging buffers (100k bars, 42 features, 4 targets)"); - crate::cuda_pipeline::GpuBufferPool::new(100_000, 42, 4) - }); - - // GPU pipeline: double-buffered loader for zero-downtime fold transitions - let double_buffer = device.is_cuda().then(|| { - crate::cuda_pipeline::double_buffer::DoubleBufferedLoader::new(device.clone()) - }); - - // Multi-GPU: auto-detect if multiple CUDA devices are available - let multi_gpu = crate::cuda_pipeline::multi_gpu::MultiGpuConfig::detect() - .unwrap_or(None); - if let Some(ref mg) = multi_gpu { - info!("Multi-GPU: {} devices detected, data parallelism enabled", mg.world_size); - } - - Ok(Self { - agent: Arc::new(RwLock::new(agent)), - hyperparams, - device, - metrics: Arc::new(RwLock::new(TrainingMetrics::new())), - loss_history: Vec::new(), - q_value_history: Vec::new(), - best_val_loss: f64::INFINITY, // Start with worst possible loss - val_data: Vec::new(), - val_loss_history: Vec::new(), - sharpe_history: Vec::new(), - best_sharpe: f64::NEG_INFINITY, // C4: Start with worst possible Sharpe - best_epoch: 0, - gradient_logging_step: 0, - collapse_warmup_buffer_size: original_buffer_size, - portfolio_tracker, - feature_stats: None, // WAVE 3 FIX #2: Start with None, collect stats in epochs 0-10 - recent_actions: VecDeque::with_capacity(100), - reward_fn, - - // WAVE 16S: Adaptive risk management - kelly_optimizer, - trade_history: VecDeque::with_capacity(500), - volatility_returns: VecDeque::with_capacity(20), // Use default instead of moved hyperparams - pnl_history: VecDeque::with_capacity(1000), - - // Wave 16 Portfolio Features - enable_action_masking, - max_position, - multi_asset_portfolio, - stress_tester, - - // Wave 16 Core Risk Features - drawdown_monitor, - position_limiter, - circuit_breaker, - - // WAVE 3.10: Microstructure feature calculators - micro_high_low_spread: HighLowSpread::default(), - micro_vw_spread: VolumeWeightedSpread::default(), - micro_tick_count: TickCount::default(), - micro_inter_arrival: InterArrivalTime::default(), - micro_buy_sell_imbalance: BuySellImbalance::default(), - micro_kyle_lambda: KyleLambda::default(), - micro_price_impact: PriceImpact::default(), - micro_variance_ratio: VarianceRatio::default(), - last_timestamp_ns: 0, - last_close: 0.0, - - // WAVE 1.1: Triple barrier integration - triple_barrier, - active_position_tracker: None, - previous_simulated_position: 0.0, // WAVE P3: Start with flat position - - // WAVE 1.2: Safety Infrastructure Integration (8 Systems) - safety_loss_history: VecDeque::with_capacity(30), - safety_loss_plateau_counter: 0, - safety_action_counts: std::collections::HashMap::new(), - safety_memory_manager: Arc::new(RwLock::new( - crate::safety::memory_manager::SafeMemoryManager::new( - &crate::safety::MLSafetyConfig::default() - ) - )), - safety_level: crate::safety::SafetyLevel::Normal, // Default to Normal mode - safety_step_counter: 0, - - feature_cache_dir: None, - - prev_epoch_q_mean: 0.0, - adaptive_tau: base_tau, - - // WAVE 24 (Agent 17): Initialize patience-based early stopping - // Use gradient_collapse_patience from hyperparams for consistency - // Set min_delta to 0.001 (0.1% improvement threshold) - early_stopping: super::early_stopping::EarlyStopping::new( - early_stopping_patience, // Reuse patience parameter (default: 5) - 0.001, // 0.1% minimum improvement - ), - - // WAVE 26 P0.6: Use pre-initialized learning rate scheduler - lr_scheduler, - - // WAVE 26 P1.8: Use pre-initialized curiosity module - curiosity_module, - - // WAVE 26 P1: Advanced DQN Features - // P1.3: Sharpe ratio reward - returns_history: VecDeque::with_capacity(sharpe_window), - sharpe_weight, - - // P1.6: Adaptive dropout - dropout_scheduler, - - // P1.7: Hindsight Experience Replay - her_buffer, - - // P1.9: Generalized Advantage Estimation - gae_calculator, - - // P1.11: Noisy sigma scheduler - noisy_sigma_scheduler, - - // WAVE 30: Structured logging integration - logging_config: LoggingConfig::default(), - metrics_aggregator: MetricsAggregator::new(), - - // WAVE 44: Multi-step returns - nstep_buffer, - - // OOM recovery: track effective batch size - current_batch_size: initial_batch_size, - - // Q-value estimation: periodic (every 50 steps) instead of every step - cached_avg_q: 0.0, - q_estimation_counter: 0, - - // GPU pipeline: pre-uploaded training data (initialized lazily at first epoch) - gpu_data: None, - #[cfg(feature = "cuda")] - gpu_portfolio_sim: None, - #[cfg(feature = "cuda")] - targets_raw_cuda: None, - #[cfg(feature = "cuda")] - features_raw_cuda: None, - #[cfg(feature = "cuda")] - gpu_experience_collector: None, - #[cfg(feature = "cuda")] - gpu_action_selector: None, - #[cfg(feature = "cuda")] - training_guard: None, - #[cfg(feature = "cuda")] - gpu_monitoring: None, - #[cfg(feature = "cuda")] - cached_n_episodes: None, - - // GPU pipeline: staging buffer pool (auto-initialized on CUDA devices) - buffer_pool, - - // GPU pipeline: double-buffered loader for fold transitions - double_buffer, - - // GPU walk-forward: initialized lazily when enable_gpu_walk_forward=true - #[cfg(feature = "cuda")] - gpu_walk_forward: None, - - // Multi-GPU: auto-detected data parallelism - multi_gpu, - - // OFI features: populated during data loading when MBP-10 data is available - ofi_features: None, - ofi_val_offset: 0, - - // Phase C: Fill simulation and smart order routing - fill_simulator: FillSimulator::default(), - vol_ema: 0.01, // Initial volatility estimate (1% daily) - median_vol: 0.01, // Slowly adapting baseline - }) - } - - /// Two-phase stress tester initialization. - /// - /// Call this after constructing `DQNTrainer` to resolve the circular dependency: - /// `DQNStressTester::new()` requires a `DQNTrainer`, so the tester cannot be - /// created *during* trainer construction. This method builds a lightweight - /// inner trainer (with stress testing itself disabled to avoid recursion) and - /// hands it to `DQNStressTester::new()`. - /// - /// No-op if `hyperparams.enable_stress_testing` is `false`. - pub fn init_stress_tester(&mut self) -> Result<()> { - if !self.hyperparams.enable_stress_testing { - return Ok(()); - } - - // Clone hyperparams with stress testing disabled so the inner trainer - // does not recursively try to initialise its own stress tester. - // Also disable GPU-heavy features that the stress tester doesn't need — - // otherwise we double the VRAM usage (3 extra regime heads, 3 extra GPU PER buffers, - // experience collector, etc.) which causes OOM on 4GB GPUs. - let mut inner_hp = self.hyperparams.clone(); - inner_hp.enable_stress_testing = false; - inner_hp.enable_regime_qnetwork = false; - inner_hp.use_per = false; - inner_hp.buffer_size = 1; - inner_hp.enable_gpu_experience_collector = false; - - let inner_trainer = Self::new_with_device(inner_hp, self.device.clone()) - .context("Failed to create inner DQNTrainer for stress tester")?; - - let tester = crate::dqn::stress_testing::DQNStressTester::new(inner_trainer)?; - self.stress_tester = Some(Arc::new(tester)); - info!("Stress testing enabled (8 scenarios)"); - - Ok(()) - } - - /// Get a reference to the double-buffered loader, if GPU is active. - pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> { - self.double_buffer.as_ref() - } - - /// Get a mutable reference to the double-buffered loader, if GPU is active. - pub fn double_buffer_mut(&mut self) -> Option<&mut crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> { - self.double_buffer.as_mut() - } - - /// Set feature cache directory for faster hyperopt - /// - /// Enables loading pre-computed features from disk instead of recomputing them - pub fn with_feature_cache(mut self, cache_dir: PathBuf) -> Self { - self.feature_cache_dir = Some(cache_dir); - self - } - - /// Train DQN on market data from DBN files - /// - /// # Arguments - /// - /// * `dbn_data_dir` - Directory containing DBN files (e.g., "test_data/real/databento/ml_training/") - /// * `checkpoint_callback` - Callback for saving checkpoints (epoch, model_data, is_final) -> `Result` - /// - /// # Returns - /// - /// Training metrics (loss, accuracy, gradient norms, Q-values) - pub async fn train( - &mut self, - dbn_data_dir: &str, - checkpoint_callback: F, - ) -> Result - where - F: FnMut(usize, Vec, bool) -> Result + Send, - { - info!( - "Starting DQN training for {} epochs with batch size {}", - self.hyperparams.epochs, self.hyperparams.batch_size - ); - - // Load market data from DBN files (ALL data for walk-forward or single-pass) - let (training_data, val_data) = self.load_training_data(dbn_data_dir).await?; - - info!( - "Loaded {} training samples, {} validation samples", - training_data.len(), - val_data.len() - ); - - // GPU walk-forward: upload ALL data to GPU, run expanding-window folds - #[cfg(feature = "cuda")] - if self.hyperparams.enable_gpu_walk_forward && self.device.is_cuda() { - // Merge train+val into a single dataset for walk-forward splitting - let mut all_data = training_data; - all_data.extend(val_data); - info!( - "GPU walk-forward enabled: {} total bars, uploading to VRAM", - all_data.len(), - ); - return self.train_walk_forward(&all_data, checkpoint_callback).await; - } - - // Standard single-pass training - self.ofi_val_offset = training_data.len(); - self.val_data = val_data; - self.train_with_data_full_loop(&training_data, checkpoint_callback) - .await - } - - /// Train with preloaded data (skips disk I/O and feature extraction). - /// - /// Accepts pre-split training and validation data that was loaded once and - /// cached across hyperopt trials. This avoids re-reading 36 `.dbn.zst` files - /// and re-extracting 42 features on every trial, eliminating minutes of GPU - /// idle time at each trial boundary. - /// - /// # Arguments - /// - /// * `training_data` - Pre-extracted (features, targets) for training split - /// * `val_data` - Pre-extracted (features, targets) for validation split - /// * `checkpoint_callback` - Checkpoint save callback - /// - /// # Returns - /// - /// Training metrics from the completed run - /// - /// # Errors - /// - /// Returns error if the training loop fails - pub async fn train_with_preloaded_data( - &mut self, - training_data: Vec<(FeatureVector, Vec)>, - val_data: Vec<(FeatureVector, Vec)>, - checkpoint_callback: F, - ) -> Result - where - F: FnMut(usize, Vec, bool) -> Result + Send, - { - info!( - "Starting DQN training with preloaded data: {} train, {} val samples", - training_data.len(), - val_data.len() - ); - - // Store validation data for loss computation - self.ofi_val_offset = training_data.len(); - self.val_data = val_data; - - // Use the common training loop (Wave 12 Group 3 refactor) - self.train_with_data_full_loop(&training_data, checkpoint_callback) - .await - } - - /// Train with shared preloaded data (zero-copy for hyperopt). - /// - /// Same as [`train_with_preloaded_data`] but accepts `Arc`-wrapped data, - /// avoiding a ~150 MB deep clone per hyperopt trial. - pub async fn train_with_shared_data( - &mut self, - training_data: &[(FeatureVector, Vec)], - val_data: Vec<(FeatureVector, Vec)>, - checkpoint_callback: F, - ) -> Result - where - F: FnMut(usize, Vec, bool) -> Result + Send, - { - info!( - "Starting DQN training with shared data: {} train, {} val samples", - training_data.len(), - val_data.len() - ); - - self.ofi_val_offset = training_data.len(); - self.val_data = val_data; - self.train_with_data_full_loop(training_data, checkpoint_callback) - .await - } - - /// Train with GPU-resident walk-forward cross-validation. - /// - /// Uploads the ENTIRE dataset to GPU VRAM once, then runs expanding-window - /// walk-forward: each fold trains on [0..T], validates on [T..V], tests on - /// [V..E]. Fold transitions are zero-copy (index range changes only). - /// - /// Returns the metrics from the LAST fold (most data, most representative). - #[cfg(feature = "cuda")] - pub async fn train_walk_forward( - &mut self, - training_data: &[(FeatureVector, Vec)], - mut checkpoint_callback: F, - ) -> Result - where - F: FnMut(usize, Vec, bool) -> Result + Send, - { - use crate::cuda_pipeline::gpu_walk_forward::{GpuWalkForwardConfig, GpuWalkForwardData}; - - let wf_config = GpuWalkForwardConfig { - initial_train_fraction: self.hyperparams.wf_initial_train_fraction, - val_fraction: self.hyperparams.wf_val_fraction, - test_fraction: self.hyperparams.wf_test_fraction, - step_fraction: self.hyperparams.wf_step_fraction, - }; - - // Upload ALL data to GPU once - let wf_data = GpuWalkForwardData::upload( - training_data, - self.ofi_features.as_deref(), - &wf_config, - &self.device, - ).map_err(|e| anyhow::anyhow!("GPU walk-forward upload: {e}"))?; - - let num_folds = wf_data.num_folds(); - if num_folds == 0 { - return Err(anyhow::anyhow!( - "Insufficient data for walk-forward: {} bars, need at least {} for one fold", - training_data.len(), - ((wf_config.initial_train_fraction + wf_config.val_fraction + wf_config.test_fraction) * training_data.len() as f64) as usize, - )); - } - - info!( - "GPU walk-forward: {} folds, {:.1} MB VRAM, {} total bars", - num_folds, wf_data.vram_bytes as f64 / 1_048_576.0, wf_data.total_bars, - ); - - // Store GPU walk-forward data and cudarc buffers for the experience collector - self.features_raw_cuda = Some(wf_data.features); - self.targets_raw_cuda = Some(wf_data.targets); - - let mut last_metrics = TrainingMetrics::new(); - - for fold_idx in 0..num_folds { - let fold = wf_data.folds.get(fold_idx).ok_or_else(|| { - anyhow::anyhow!("Fold {fold_idx} out of range") - })?; - - info!( - "=== Walk-Forward Fold {}/{} === train: {} bars, val: {} bars, test: {} bars", - fold_idx + 1, num_folds, fold.train_len(), fold.val_len(), fold.test_len(), - ); - - // Split training_data into fold's train and val slices (for CPU-side data) - let fold_train = &training_data[fold.train_start..fold.train_end]; - let fold_val: Vec<(FeatureVector, Vec)> = - training_data[fold.val_start..fold.val_end].to_vec(); - - // Store validation data for this fold - self.ofi_val_offset = fold.train_end; - self.val_data = fold_val; - - // Reset training state for new fold - self.gpu_data = None; // Force re-upload via DqnGpuData for the fold's range - self.best_sharpe = f64::NEG_INFINITY; - self.best_val_loss = f64::INFINITY; - self.loss_history.clear(); - self.q_value_history.clear(); - self.val_loss_history.clear(); - self.sharpe_history.clear(); - - // Run training loop on this fold's data - last_metrics = self - .train_with_data_full_loop(fold_train, &mut checkpoint_callback) - .await?; - - info!( - "Fold {}/{} complete: loss={:.6}, epochs={}", - fold_idx + 1, num_folds, - last_metrics.loss, - last_metrics.epochs_trained, - ); - } - - // Clean up GPU walk-forward buffers (features/targets already stored in self) - self.gpu_walk_forward = None; - - Ok(last_metrics) - } - - /// Full training loop with existing logic (Wave 12 Group 3) - /// Calculate average metrics for an epoch - fn calculate_epoch_metrics( - epoch_loss: f64, - epoch_q_value: f64, - epoch_gradient_norm: f64, - samples_processed: usize, - ) -> (f64, f64, f64) { - if samples_processed > 0 { - let count = samples_processed as f64; - ( - epoch_loss / count, - epoch_q_value / count, - epoch_gradient_norm / count, - ) - } else { - (0.0, 0.0, 0.0) - } - } - - /// Compute validation loss on held-out data - /// WAVE 10.6: Batched validation for 5-10x speedup - - /// Collect Q-value statistics from replay buffer - /// - /// Samples experiences from the replay buffer and computes Q-value statistics - /// (min, max, mean, std) for adaptive C51 bounds calculation. - /// - /// # Returns - /// - /// QValueStats with min/max/mean/std of Q-values - async fn collect_qvalue_statistics(&self) -> Result { - let agent = self.agent.read().await; - - // Determine sample size (min of buffer size or 1000) - let buffer_size = agent.get_replay_buffer_size()?; - let sample_size = buffer_size.min(1000); - - if sample_size == 0 { - return Err(crate::MLError::TrainingError( - "Replay buffer is empty, cannot collect Q-value statistics".to_owned() - )); - } - - // Sample experiences from replay buffer - let batch_sample = agent.memory().sample(sample_size)?; - let state_dim = agent.get_state_dim(); - - // GPU PER path: use gpu_batch.states directly (experiences vec is empty) - #[allow(unused_assignments, unused_mut)] - let mut batch_tensor_opt: Option = None; - #[cfg(feature = "cuda")] - { - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor_opt = Some( - gpu.states.to_dtype(training_dtype(agent.device())) - .map_err(|e| crate::MLError::ModelError(format!("GPU Q-stat states dtype cast: {}", e)))? - ); - } - } - let batch_tensor = if let Some(t) = batch_tensor_opt { - t - } else { - let states: Vec = batch_sample.experiences - .iter() - .flat_map(|exp| exp.state.iter().copied()) - .collect(); - Tensor::from_vec(states, (sample_size, state_dim), agent.device()) .map_err(|e| crate::MLError::ModelError(format!("Failed to create batch tensor: {}", e)))? - .to_dtype(training_dtype(agent.device())) - .map_err(|e| crate::MLError::ModelError(format!("Failed to cast batch tensor to training dtype: {}", e)))? - }; - - // Forward pass to get Q-values [batch_size, num_actions] - let q_values = agent.forward(&batch_tensor)?; - - // GPU-side statistics: flatten Q-values and compute min/max/mean/std on device. - // Only 4 scalar readbacks (16 bytes) instead of downloading the entire tensor. - let q_f32 = q_values - .to_dtype(candle_core::DType::F32) - .map_err(|e| crate::MLError::ModelError(format!("Q-value F32 cast: {}", e)))?; - let q_flat = q_f32 - .flatten_all() - .map_err(|e| crate::MLError::ModelError(format!("Q-value flatten: {}", e)))?; - let count = q_flat.elem_count(); - - let min = q_flat.min(0) - .and_then(|t| t.to_vec0::()) - .map_err(|e| crate::MLError::ModelError(format!("Q-value min: {}", e)))? as f64; - let max = q_flat.max(0) - .and_then(|t| t.to_vec0::()) - .map_err(|e| crate::MLError::ModelError(format!("Q-value max: {}", e)))? as f64; - let mean_scalar = q_flat.mean_all() - .and_then(|t| t.to_vec0::()) - .map_err(|e| crate::MLError::ModelError(format!("Q-value mean: {}", e)))?; - let mean = mean_scalar as f64; - - // std = sqrt(mean((x - mean)^2)) — all on device - let mean_tensor = q_flat.mean_all() - .map_err(|e| crate::MLError::ModelError(format!("Q-value mean tensor: {}", e)))?; - let variance = q_flat.broadcast_sub(&mean_tensor) - .and_then(|d| d.sqr()) - .and_then(|sq| sq.mean_all()) - .and_then(|v| v.to_vec0::()) - .map_err(|e| crate::MLError::ModelError(format!("Q-value variance: {}", e)))? as f64; - let std = variance.sqrt(); - - Ok(QValueStats { - min, - max, - mean, - std, - sample_count: count, - }) - } - - /// Calculate adaptive bounds with margin - /// - /// # Arguments - /// - /// * `stats` - Q-value statistics from Phase 1 - /// * `margin` - Safety margin as fraction (e.g., 0.3 = 30%) - /// - /// # Returns - /// - /// Tuple of (v_min, v_max) with safety margin applied - fn calculate_adaptive_bounds(stats: &QValueStats, margin: f64) -> (f64, f64) { - let range = stats.max - stats.min; - let v_min = stats.min - range * margin; - let v_max = stats.max + range * margin; - // Cap at ±10,000 to prevent explosion - (v_min.max(-10000.0), v_max.min(10000.0)) - } - - /// Reinitialize categorical distribution with new bounds - async fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), crate::MLError> { - let mut agent = self.agent.write().await; - match &mut *agent { - DQNAgentType::Standard(agent) => agent.reinit_categorical_distribution(v_min, v_max)?, - DQNAgentType::RegimeConditional(agent) => agent.reinit_categorical_distribution(v_min, v_max)?, - } - Ok(()) - } - - #[allow(unused_variables, unreachable_code, unused_mut)] - async fn compute_validation_loss(&mut self) -> Result { - if self.val_data.is_empty() { - return Ok(0.0); - } - - // Save current epsilon and force to 0 for deterministic evaluation - let original_epsilon = self.get_epsilon().await?; - self.set_epsilon(0.0).await?; // Pure greedy selection - - let sample_size = self.val_data.len().min(1000); // Sample up to 1000 for speed - - // P4: Get aligned state_dim from the agent (already tensor-core aligned at construction) - // so we can pre-allocate the flat GPU buffer without an intermediate Vec>. - let aligned_state_dim = { - let agent = self.agent.read().await; - agent.get_state_dim() - }; - - // P4: Pre-allocate a single zero-initialized flat buffer for the entire batch. - // Zero-init handles tensor core padding columns automatically -- no separate - // padding branch needed. Each state is written at offset `i * aligned_state_dim`. - // This eliminates N intermediate Vec allocations (one per sample) and - // the subsequent flat_map copy pass from the WAVE 10.6 implementation. - let mut batched_states: Vec = vec![0.0_f32; sample_size * aligned_state_dim]; - let mut states = Vec::with_capacity(sample_size); - let mut next_states = Vec::with_capacity(sample_size); - let mut actions_for_rewards = Vec::with_capacity(sample_size); - // GPU reward path: collect close prices for PnL-based Sharpe on device - let mut val_current_closes: Vec = Vec::with_capacity(sample_size); - let mut val_next_closes: Vec = Vec::with_capacity(sample_size); - - for (i, (feature_vec, target)) in self.val_data.iter().take(sample_size).enumerate() { - let current_close = if target.len() >= 2 { - target[0] - } else { - feature_vec[3] - }; - let next_close = if target.len() >= 2 { - target[1] - } else { - current_close - }; - let close_price = rust_decimal::Decimal::try_from(current_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - // OFI fix: pass global OFI index so validation gets real OFI features - let ofi_idx = self.ofi_val_offset + i; - let state = self.feature_vector_to_state_with_ofi(feature_vec, Some(close_price), Some(ofi_idx))?; - - let next_close_price = - rust_decimal::Decimal::try_from(next_close).unwrap_or(rust_decimal::Decimal::ZERO); - let next_state = self.feature_vector_to_state_with_ofi(feature_vec, Some(next_close_price), Some(ofi_idx))?; - - // P4: Write state vector directly into the flat buffer at the correct row offset. - // Raw dims (e.g. 45 or 53) are shorter than aligned (48 or 56); trailing - // positions stay zero from the vec![0.0; ..] init -- no explicit pad needed. - let sv = state.to_vector(); - let row_start = i * aligned_state_dim; - let copy_len = sv.len().min(aligned_state_dim); - let buf_len = batched_states.len(); - let dst = batched_states.get_mut(row_start..row_start + copy_len) - .ok_or_else(|| anyhow::anyhow!( - "Validation buffer overrun: row_start={}, copy_len={}, buf_len={}", - row_start, copy_len, buf_len - ))?; - let src = sv.get(..copy_len).ok_or_else(|| anyhow::anyhow!( - "State vector shorter than expected: len={}, copy_len={}", - sv.len(), copy_len - ))?; - dst.copy_from_slice(src); - - states.push(state); - next_states.push(next_state); - val_current_closes.push(current_close as f32); - val_next_closes.push(next_close as f32); - } - - // P4: Single batched forward pass -- tensor created directly from the flat buffer, - // no intermediate Vec> needed. - let agent = self.agent.read().await; - let batch_tensor = Tensor::from_vec(batched_states, (sample_size, aligned_state_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched validation tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast validation tensor to training dtype: {}", e))?; - - let batch_q_values = agent.forward(&batch_tensor) - .map_err(|e| anyhow::anyhow!("Batched validation forward pass failed: {}", e))?; - - // Get branching Q-values if branching is enabled - let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = - if self.hyperparams.use_branching { - agent - .batch_branching_q_values(&batch_tensor) - .map_err(|e| anyhow::anyhow!("Validation branching Q-values failed: {e}"))? - } else { - None - }; - - drop(agent); // Release lock early - - #[cfg(not(feature = "cuda"))] - return Err(anyhow::anyhow!("Validation requires CUDA — enable the `cuda` feature")); - - // Fused GPU greedy action selection (epsilon=0.0) + GPU routing. - #[cfg(feature = "cuda")] - { - if self.gpu_action_selector.is_none() && self.device.is_cuda() { - let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( - &self.device, - self.hyperparams.batch_size.max(sample_size).max(8192), - 0xDEAD_BEEF_CAFE_u64, - ).map_err(|e| anyhow::anyhow!("Validation GPU action selector init failed: {e}"))?; - self.gpu_action_selector = Some(selector); - } - - let selector = self.gpu_action_selector.as_mut() - .ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?; - - let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { - selector - .select_actions_branching(q_exp, q_ord, q_urg, 0.0) - .map_err(|e| anyhow::anyhow!("Validation GPU branching select failed: {e}"))? - } else { - let exposure_tensor = selector - .select_actions(&batch_q_values, 0.0, sample_size, 5) - .map_err(|e| anyhow::anyhow!("Validation GPU greedy select failed: {e}"))?; - selector.route_exposure_to_factored( - &exposure_tensor, sample_size, - self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, - self.vol_ema as f32, self.median_vol as f32, - ).map_err(|e| anyhow::anyhow!("Validation GPU route failed: {e}"))? - }; - - // --- GPU PnL-based reward + Sharpe computation --- - // Instead of reading factored indices back to CPU and looping through - // the scalar reward function, compute PnL rewards entirely on GPU: - // reward_i = direction_i * (next_close_i - current_close_i) / current_close_i - // Then compute Sharpe = mean(rewards) / std(rewards) * sqrt(252) on device - // with a single scalar readback for the final value. - - // 1) Build close-price tensors on GPU - let current_closes_t = Tensor::from_vec( - val_current_closes, &[sample_size], &self.device, - ).map_err(|e| anyhow::anyhow!("GPU val current_closes tensor: {e}"))?; - let next_closes_t = Tensor::from_vec( - val_next_closes, &[sample_size], &self.device, - ).map_err(|e| anyhow::anyhow!("GPU val next_closes tensor: {e}"))?; - - // 2) Extract exposure index from factored tensor: - // factored_index = exposure * 9 + order * 3 + urgency - // => exposure_idx = factored_index / 9 - let factored_f32 = factored_tensor - .to_dtype(candle_core::DType::F32) - .map_err(|e| anyhow::anyhow!("GPU val factored to f32: {e}"))?; - let nine = Tensor::new(&[9.0_f32], &self.device) - .map_err(|e| anyhow::anyhow!("GPU val nine const: {e}"))?; - let exposure_idx_f32 = factored_f32 - .broadcast_div(&nine) - .map_err(|e| anyhow::anyhow!("GPU val exposure div: {e}"))? - .floor() - .map_err(|e| anyhow::anyhow!("GPU val exposure floor: {e}"))?; - - // 3) Map exposure index to direction multiplier via lookup table: - // [0=Short100→-1.0, 1=Short50→-0.5, 2=Flat→0.0, 3=Long50→0.5, 4=Long100→1.0] - let direction_lut = Tensor::new( - &[-1.0_f32, -0.5, 0.0, 0.5, 1.0], &self.device, - ).map_err(|e| anyhow::anyhow!("GPU val direction LUT: {e}"))?; - let exposure_idx_u32 = exposure_idx_f32 - .to_dtype(candle_core::DType::U32) - .map_err(|e| anyhow::anyhow!("GPU val exposure to u32: {e}"))?; - let directions = direction_lut - .index_select(&exposure_idx_u32, 0) - .map_err(|e| anyhow::anyhow!("GPU val direction gather: {e}"))?; - - // 4) Compute PnL-based rewards: direction * (next - current) / current - let price_returns = next_closes_t - .sub(¤t_closes_t) - .map_err(|e| anyhow::anyhow!("GPU val price diff: {e}"))? - .broadcast_div(¤t_closes_t) - .map_err(|e| anyhow::anyhow!("GPU val price returns div: {e}"))?; - let rewards = directions - .mul(&price_returns) - .map_err(|e| anyhow::anyhow!("GPU val rewards mul: {e}"))?; - - // 5) Compute Sharpe on GPU: mean / std * sqrt(252) - let mean_t = rewards.mean_all() - .map_err(|e| anyhow::anyhow!("GPU val rewards mean: {e}"))?; - let mean_scalar: f64 = mean_t - .to_scalar::() - .map_err(|e| anyhow::anyhow!("GPU val mean readback: {e}"))? as f64; - let var_t = rewards - .broadcast_sub(&mean_t) - .map_err(|e| anyhow::anyhow!("GPU val rewards center: {e}"))? - .sqr() - .map_err(|e| anyhow::anyhow!("GPU val rewards sqr: {e}"))? - .mean_all() - .map_err(|e| anyhow::anyhow!("GPU val rewards var: {e}"))?; - let var_scalar: f64 = var_t - .to_scalar::() - .map_err(|e| anyhow::anyhow!("GPU val var readback: {e}"))? as f64; - - let std_val = var_scalar.sqrt(); - let val_sharpe = if std_val > 1e-10 { - (mean_scalar / std_val) * (252.0_f64).sqrt() - } else { - 0.0 - }; - - // Restore original epsilon after evaluation - self.set_epsilon(original_epsilon).await?; - - // Return negative Sharpe as the "loss" (lower = better Sharpe) - return Ok(-val_sharpe); - } - - // Fallback: CPU reward path (only reachable if cuda cfg gate above is bypassed) - #[allow(unreachable_code)] - { - let mut val_rewards = Vec::with_capacity(sample_size); - let recent_actions_vec: Vec = - self.recent_actions.iter().copied().collect(); - - for (i, action) in actions_for_rewards.iter().enumerate() { - let state_ref = states.get(i).ok_or_else(|| { - anyhow::anyhow!("Validation state index {} out of bounds (len={})", i, states.len()) - })?; - let next_ref = next_states.get(i).ok_or_else(|| { - anyhow::anyhow!("Validation next_state index {} out of bounds (len={})", i, next_states.len()) - })?; - let reward_decimal = self.reward_fn.calculate_reward( - *action, - state_ref, - next_ref, - &recent_actions_vec, - )?; - val_rewards.push(reward_decimal.to_f32().unwrap_or(0.0) as f64); - } - - self.set_epsilon(original_epsilon).await?; - - let n = val_rewards.len() as f64; - let mean = val_rewards.iter().sum::() / n; - let variance = val_rewards.iter().map(|r| (r - mean).powi(2)).sum::() / n; - let std = variance.sqrt(); - let val_sharpe = if std > 1e-10 { - (mean / std) * (252.0_f64).sqrt() - } else { - 0.0 - }; - - Ok(-val_sharpe) - } - } - - /// Get Q-values for a given state - async fn get_q_values(&self, state: &TradingState) -> Result> { - let agent = self.agent.read().await; - let state_vec = state.to_vector(); - let raw_dim = state_vec.len(); - let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - let padded: Vec = if aligned > raw_dim { - let mut v = state_vec.to_vec(); - v.resize(aligned, 0.0); - v - } else { - state_vec.to_vec() - }; - let state_tensor = Tensor::new(&*padded, &self.device)?.unsqueeze(0)?; // Add batch dimension - - let q_values_tensor = agent.forward(&state_tensor)?.squeeze(0)?; - let n = q_values_tensor.elem_count(); - let mut q_values = Vec::with_capacity(n); - for i in 0..n { - q_values.push( - q_values_tensor.narrow(0, i, 1)?.to_scalar::()? as f64, - ); - } - Ok(q_values) - } - - /// Check if early stopping criteria are met - fn check_early_stopping(&self, avg_q_value: f64, epoch: usize) -> Option { - if !self.hyperparams.early_stopping_enabled - || epoch + 1 < self.hyperparams.min_epochs_before_stopping - { - return None; - } - - // Criterion 1: Q-value floor check - if avg_q_value < self.hyperparams.q_value_floor { - return Some(format!( - "Q-value {:.4} below floor threshold {:.4}", - avg_q_value, self.hyperparams.q_value_floor - )); - } - - // C4 FIX: Criterion 2 — Sharpe plateau check (was val-loss plateau). - // Sharpe directly measures trading quality. A plateau means the model - // has stopped improving its trading strategy, even if TD-loss still moves. - if self.sharpe_history.len() >= self.hyperparams.plateau_window { - let window = self.hyperparams.plateau_window; - let recent_sharpes: Vec = self - .sharpe_history - .iter() - .rev() - .take(window) - .copied() - .collect(); - - if let (Some(&newest), Some(&oldest)) = (recent_sharpes.first(), recent_sharpes.last()) { - // For Sharpe, improvement = newest - oldest (higher is better) - let improvement = newest - oldest; - - if improvement < 0.01 { - let msg = if improvement < -0.01 { - format!( - "Sharpe worsening detected (delta: {:.4}, window: {})", - improvement, - window - ) - } else { - format!( - "Sharpe plateau detected (improvement: {:.4}, window: {})", - improvement, - window - ) - }; - return Some(msg); - } - } - } - - None - } - - /// Create final training metrics - async fn create_final_metrics( - &self, - total_loss: f64, - total_q_value: f64, - total_gradient_norm: f64, - total_reward: f64, - num_epochs: usize, - training_duration: std::time::Duration, - early_stopped: bool, - total_action_counts: [usize; 5], // 5 exposure levels - total_factored_action_counts: [usize; 45], // 45 factored actions - ) -> Result { - let final_loss = total_loss / num_epochs as f64; - let avg_q_value_final = total_q_value / num_epochs as f64; - let avg_grad_norm_final = total_gradient_norm / num_epochs as f64; - let avg_episode_reward = total_reward / num_epochs as f64; - - let mut metrics = TrainingMetrics { - loss: final_loss, - accuracy: 0.0, - precision: 0.0, - recall: 0.0, - f1_score: 0.0, - training_time_seconds: training_duration.as_secs_f64(), - epochs_trained: num_epochs as u32, - convergence_achieved: final_loss < 1.0, - additional_metrics: std::collections::HashMap::new(), - }; - - metrics.add_metric("avg_q_value", avg_q_value_final); - metrics.add_metric("avg_gradient_norm", avg_grad_norm_final); - metrics.add_metric("final_epsilon", self.get_epsilon().await.unwrap_or(0.1)); - metrics.add_metric("avg_episode_reward", avg_episode_reward); - - // Action metrics: use factored 45-action space if branching, else 5 exposure levels - let total_factored: usize = total_factored_action_counts.iter().sum(); - let total_exposure: usize = total_action_counts.iter().sum(); - let total_actions = total_factored.max(total_exposure); - if total_actions > 0 { - // Factored 45-action diversity (primary metric when branching) - if total_factored > 0 { - let unique_factored = total_factored_action_counts.iter() - .filter(|&&count| count > 0).count(); - let factored_diversity = (unique_factored as f64 / 45.0) * 100.0; - metrics.add_metric("action_diversity", factored_diversity); - metrics.add_metric("factored_unique_actions", unique_factored as f64); - metrics.add_metric("action_space_size", 45.0); - - // Active factored actions (used >0.5% of the time) - let active_threshold = (total_factored as f64 * 0.005).max(1.0); - let active_count = total_factored_action_counts.iter() - .filter(|&&count| count as f64 >= active_threshold) - .count(); - let active_diversity_pct = (active_count as f64 / 45.0) * 100.0; - metrics.add_metric("active_actions_count", active_count as f64); - metrics.add_metric("active_diversity_pct", active_diversity_pct); - - // Top factored actions - let mut sorted_actions: Vec<(usize, usize)> = total_factored_action_counts.iter() - .enumerate() - .map(|(idx, &count)| (idx, count)) - .collect(); - sorted_actions.sort_by(|a, b| b.1.cmp(&a.1)); - - if let Some((top1_idx, top1_count)) = sorted_actions.first() { - let top1_pct = (*top1_count as f64 / total_factored as f64) * 100.0; - metrics.add_metric("top1_action_idx", *top1_idx as f64); - metrics.add_metric("top1_action_count", *top1_count as f64); - metrics.add_metric("top1_action_pct", top1_pct); - } - - let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum(); - let top5_coverage_pct = (top5_count as f64 / total_factored as f64) * 100.0; - metrics.add_metric("top5_coverage_pct", top5_coverage_pct); - } else { - // Fallback: 5 exposure levels - let unique_actions = total_action_counts.iter() - .filter(|&&count| count > 0).count(); - let action_diversity = (unique_actions as f64 / 5.0) * 100.0; - metrics.add_metric("action_diversity", action_diversity); - metrics.add_metric("action_space_size", 5.0); - - let active_threshold = (total_exposure as f64 * 0.005).max(1.0); - let active_count = total_action_counts.iter() - .filter(|&&count| count as f64 >= active_threshold) - .count(); - let active_diversity_pct = (active_count as f64 / 5.0) * 100.0; - metrics.add_metric("active_actions_count", active_count as f64); - metrics.add_metric("active_diversity_pct", active_diversity_pct); - - let mut sorted_actions: Vec<(usize, usize)> = total_action_counts.iter() - .enumerate() - .map(|(idx, &count)| (idx, count)) - .collect(); - sorted_actions.sort_by(|a, b| b.1.cmp(&a.1)); - - if let Some((top1_idx, top1_count)) = sorted_actions.first() { - let top1_pct = (*top1_count as f64 / total_exposure as f64) * 100.0; - metrics.add_metric("top1_action_idx", *top1_idx as f64); - metrics.add_metric("top1_action_count", *top1_count as f64); - metrics.add_metric("top1_action_pct", top1_pct); - } - - let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum(); - let top5_coverage_pct = (top5_count as f64 / total_exposure as f64) * 100.0; - metrics.add_metric("top5_coverage_pct", top5_coverage_pct); - } - - metrics.add_metric("total_actions", total_actions as f64); - - // Buy/sell/hold always from 5-exposure space (meaningful for P&L) - // 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100 - let sell_count: usize = total_action_counts.get(0).copied().unwrap_or(0) - + total_action_counts.get(1).copied().unwrap_or(0); - let hold_count: usize = total_action_counts.get(2).copied().unwrap_or(0); - let buy_count: usize = total_action_counts.get(3).copied().unwrap_or(0) - + total_action_counts.get(4).copied().unwrap_or(0); - metrics.add_metric("buy_count", buy_count as f64); - metrics.add_metric("sell_count", sell_count as f64); - metrics.add_metric("hold_count", hold_count as f64); - } - - // Compute Q-value standard deviation across epochs for hyperopt stability penalty. - // self.q_value_history stores per-epoch average Q-values; their std measures - // how much Q-values fluctuate during training (volatility indicator). - if self.q_value_history.len() >= 2 { - let n = self.q_value_history.len() as f64; - let mean = self.q_value_history.iter().sum::() / n; - let variance = self - .q_value_history - .iter() - .map(|&q| (q - mean).powi(2)) - .sum::() - / n; - let std_dev = variance.sqrt(); - metrics.add_metric("q_value_std", std_dev); - } else { - // Not enough data points to compute std; default to 0.0 (no volatility signal) - metrics.add_metric("q_value_std", 0.0); - } - - if early_stopped { - metrics.add_metric("early_stopped", 1.0); - } - - Ok(metrics) - } - - pub(crate) async fn train_with_data_full_loop( - &mut self, - training_data: &[(FeatureVector, Vec)], - mut checkpoint_callback: F, - ) -> Result - where - F: FnMut(usize, Vec, bool) -> Result + Send, - { - let start_time = std::time::Instant::now(); - let mut total_loss = 0.0; - let mut total_q_value = 0.0; - let mut total_gradient_norm = 0.0; - let mut total_reward = 0.0; // Track cumulative rewards across all epochs - let mut total_action_counts = [0_usize; 5]; // 5 exposure levels - let mut total_factored_action_counts = [0_usize; 45]; // 45 factored actions - - // WAVE 16 (Agent 36): Log target update strategy (one-time at training start) - match self.hyperparams.target_update_mode { - TargetUpdateMode::Soft => { - let half_life = convergence_half_life(self.hyperparams.tau); - info!("🎯 WAVE 16: Using soft target updates (Polyak averaging)"); - info!(" • Tau: {}", self.hyperparams.tau); - info!(" • Convergence half-life: {} steps", half_life as usize); - info!(" • Strategy: Smooth Q-value tracking (50-70% variance reduction)"); - }, - TargetUpdateMode::Hard => { - info!("⚠️ WAVE 16: Using hard target updates (legacy mode)"); - info!(" • Update frequency: every 1000 steps"); - info!(" • Warning: Sudden Q-value shifts may cause instability"); - }, - } - - // Rainbow DQN Component Status (one-time at training start) - info!("🌈 Rainbow DQN Components:"); - info!(" ✅ Double DQN (always enabled)"); - - if self.hyperparams.use_dueling { - info!(" ✅ Dueling Networks (value/advantage streams, hidden_dim={})", self.hyperparams.dueling_hidden_dim); - } else { - info!(" ❌ Dueling Networks"); - } - - if self.hyperparams.use_per { - info!(" ✅ Prioritized Experience Replay (α={}, β={}→1.0)", - self.hyperparams.per_alpha, self.hyperparams.per_beta_start); - } else { - info!(" ❌ Prioritized Experience Replay"); - } - - if self.hyperparams.n_steps > 1 { - info!(" ✅ N-Step Returns (n={})", self.hyperparams.n_steps); - } else { - info!(" ❌ N-Step Returns (n=1, standard TD)"); - } - - if self.hyperparams.use_distributional { - info!(" ✅ Categorical DQN (atoms={}, V=[{}, {}])", - self.hyperparams.num_atoms, self.hyperparams.v_min, self.hyperparams.v_max); - } else { - info!(" ❌ Categorical DQN / C51"); - } - - if self.hyperparams.use_noisy_nets { - info!(" ✅ Noisy Networks (σ_init={})", self.hyperparams.noisy_sigma_init); - } else { - info!(" ❌ Noisy Networks"); - } - - // When noisy nets are enabled, epsilon-greedy is replaced by noisy_epsilon_floor. - // Set stored epsilon to the floor value so get_epsilon() reports it accurately. - // select_action() reads noisy_epsilon_floor from DQNConfig directly. - if self.hyperparams.use_noisy_nets { - let floor = self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05); - let mut agent = self.agent.write().await; - agent.set_epsilon(floor); - info!(" Noisy nets active: epsilon set to noisy_epsilon_floor={:.4}", floor); - } - - // Training loop - for epoch in 0..self.hyperparams.epochs { - // WAVE 30: Reset metrics aggregator for new epoch - self.metrics_aggregator.reset(); - - // Phase B1: Clear pnl_history per epoch so Sharpe reflects THIS epoch only. - // Without this, pnl_history accumulates across all epochs (ring buffer never clears), - // causing Sharpe to freeze after the first few epochs fill the buffer. - self.pnl_history.clear(); - - // Reset reward function epoch state (DSR EMA, returns buffer, normalizer) - // for reward stationarity between epochs. - self.reward_fn.reset_epoch_state(); - - // When DSR is active, reset portfolio tracker to initial capital. - // DSR handles reward normalization internally, so compounding is unnecessary. - // Without DSR, keep compounding (Bug #15 fix — reset kills learning signal). - if self.hyperparams.use_dsr { - self.portfolio_tracker.reset(); - } - - // GPU-persistent epoch state: set reset flags instead of CPU state mutation. - // Bit 0 = reset portfolio, bit 1 = reset DSR normalizer. - // Vol EMA (bit 2) is intentionally never reset between epochs (continuous tracking). - // Flags are consumed by the next kernel launch and auto-cleared. - #[cfg(feature = "cuda")] - if let Some(ref mut collector) = self.gpu_experience_collector { - let mut flags: u32 = 0; - if self.hyperparams.use_dsr { - flags |= 1; // reset portfolio - flags |= 2; // reset DSR normalizer - } - collector.set_reset_flags(flags); - } - - // WAVE 30: Log epoch start - log_epoch_start(epoch + 1, self.hyperparams.epochs, self.hyperparams.learning_rate); - - // Emit epoch gauge immediately so Prometheus/Grafana template variables resolve - // before the first epoch completes (foxhunt_training_step alone isn't enough - // because the dashboard model/fold dropdowns historically use current_epoch). - training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64); - - // Create monitor for this epoch - let mut monitor = TrainingMonitor::new(epoch + 1); - - // BUG #15 FIX (Wave 16S-V15): Portfolio compounding across epochs - // - // REMOVED: self.portfolio_tracker.reset(); - // - // Root Cause: Resetting portfolio to initial capital ($100k default) at the START - // of every epoch caused catastrophic learning signal collapse: - // - // BEFORE (BROKEN - Bug #15): - // Epoch 1: $100k → $105k (reward = +5.0%, variance = 0.0001) - // Epoch 2: $100k → $104k (reward = +4.0%, variance = 0.0001) ← RESET! - // Epoch 3: $100k → $106k (reward = +6.0%, variance = 0.0001) ← RESET! - // Result: Constant rewards ~0.004 ± 0.0001 (ZERO learning signal) - // - // AFTER (FIXED - Compounding): - // Epoch 1: $100k → $105k (reward = +5.0%, reward_std = 0.02) - // Epoch 2: $105k → $110k (reward = +5.0%, reward_std = 0.03) - // Epoch 3: $110k → $116k (reward = +5.5%, reward_std = 0.04) - // Epoch 100: $500k → $550k (reward = +10.0%, reward_std = 0.50) - // Result: Increasing reward variance (10x-100x improvement in learning signal!) - // - // Why This Matters: - // - DQN learns by observing reward differences across states/actions - // - Constant rewards (0.004 ± 0.0001) provide ZERO differentiation - // - Compounding creates natural variance: profitable strategies compound faster - // - Higher portfolio values amplify good/bad decisions (better signal-to-noise) - // - Epoch 100 reward = 10x Epoch 1 reward (massive learning signal improvement) - // - // Portfolio is initialized ONCE at trainer creation (line 600-604): - // let portfolio_tracker = PortfolioTracker::new( - // hyperparams.initial_capital, // Default: $100k - // 0.0001, // 1 basis point spread - // hyperparams.cash_reserve_percent, - // ); - // - // Portfolio only resets when starting a NEW training run (new DQNTrainer instance). - // Within a single training run, portfolio compounds across ALL epochs. - // - // However, we DO reset the drawdown high-water mark at each epoch boundary - // so the circuit breaker (>20% drawdown) only measures within-epoch drawdown. - // Without this, a compounding portfolio that hit a drawdown in epoch N would - // permanently lock out all trades in epochs N+1, N+2, ... (constant zero rewards). - self.portfolio_tracker.reset_drawdown_tracking(); - - let epoch_start = std::time::Instant::now(); - - // Phase 1: Pre-upload training data to GPU (once, reused across epochs) - if self.gpu_data.is_none() { - // Check double-buffer: skip upload if active slot already populated - let skip_upload = self.double_buffer.as_ref().is_some_and(|db| db.active().is_some()); - - if skip_upload { - info!("DoubleBuffer: active slot populated, skipping re-upload"); - } else { - let upload_result = if let Some(ref mut pool) = self.buffer_pool { - info!("GpuBufferPool: reusing pre-allocated staging buffers for {} bars", training_data.len()); - pool.upload_dqn(&training_data, &self.device) - } else { - DqnGpuData::upload(&training_data, &self.device) - }; - match upload_result { - Ok(mut gpu_data) => { - info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)", - gpu_data.num_bars, - gpu_data.feature_dim, - (gpu_data.num_bars * (42 + 4) * 4) as f64 / 1_048_576.0 - ); - // Upload OFI features to GPU if available - if let Some(ref ofi) = self.ofi_features { - match gpu_data.upload_ofi(ofi, &self.device) { - Ok(()) => info!("GPU OFI features uploaded: {} bars x 8 dims", ofi.len()), - Err(e) => { - return Err(anyhow::anyhow!("GPU OFI upload FAILED (no CPU fallback): {e}")); - } - } - } - // Set tensor core alignment so build_*_states pads output - let ofi_enabled = self.ofi_features.is_some(); - let raw_dim = if ofi_enabled { 53 } else { 45 }; - let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - gpu_data.set_aligned_state_dim(aligned_dim); - self.gpu_data = Some(gpu_data); - } - Err(e) => { - return Err(anyhow::anyhow!("GPU data pre-upload FAILED (no CPU fallback): {e}")); - } - } - } - } - - // Phase 1b: Upload raw targets + features via cudarc + init GPU portfolio sim (once) - #[cfg(feature = "cuda")] - if self.targets_raw_cuda.is_none() { - if let candle_core::Device::Cuda(ref cuda_dev) = self.device { - let stream = cuda_dev.cuda_stream(); - let target_dim = 4; - let feature_dim = 42; - let num_bars = training_data.len(); - - // Build flat targets (same as DqnGpuData::upload but for cudarc) - let mut flat_targets = Vec::with_capacity(num_bars * target_dim); - for (_, targets) in training_data { - for i in 0..target_dim { - flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32); - } - } - - match stream.memcpy_stod(&flat_targets) { - Ok(buf) => { - info!("CUDA targets_raw uploaded: {} bars × 4 ({:.1} KB)", - num_bars, (num_bars * target_dim * 4) as f64 / 1024.0); - self.targets_raw_cuda = Some(buf); - } - Err(e) => { - return Err(anyhow::anyhow!("CUDA targets_raw upload FAILED (no CPU fallback): {e}")); - } - } - - // Build flat features [num_bars * 42] for GPU experience kernel - if self.features_raw_cuda.is_none() { - let mut flat_features = Vec::with_capacity(num_bars * feature_dim); - for (features, _) in training_data { - for &v in features.iter() { - flat_features.push(v as f32); - } - } - match stream.memcpy_stod(&flat_features) { - Ok(buf) => { - info!("CUDA features_raw uploaded: {} bars × {} ({:.1} KB)", - num_bars, feature_dim, (num_bars * feature_dim * 4) as f64 / 1024.0); - self.features_raw_cuda = Some(buf); - } - Err(e) => { - return Err(anyhow::anyhow!("CUDA features_raw upload FAILED (no CPU fallback): {e}")); - } - } - } - - // Initialize GPU portfolio simulator - if self.gpu_portfolio_sim.is_none() { - if let Some(ref _targets_buf) = self.targets_raw_cuda { - use crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator; - match GpuPortfolioSimulator::new( - stream, - self.hyperparams.initial_capital as f32, - self.hyperparams.avg_spread as f32, - self.hyperparams.cash_reserve_percent as f32, - self.max_position as f32, - EPISODE_LENGTH, - training_data.len(), - ) { - Ok(sim) => { - info!("GPU portfolio simulator initialized"); - self.gpu_portfolio_sim = Some(sim); - } - Err(e) => { - return Err(anyhow::anyhow!("GPU portfolio sim init FAILED (no CPU fallback): {e}")); - } - } - } - } - } - } - - // Phase 1c: Initialize zero-roundtrip GPU experience collector (once) - // RegimeConditional agents now supported — GPU tensor batches are inserted - // into all head buffers, regime routing happens at train time via GPU masks. - // Branching DQN: the fused CUDA kernel now supports branching mode via - // use_branching flag + 16 extra weight pointers for order/urgency heads. - #[cfg(feature = "cuda")] - if self.hyperparams.enable_gpu_experience_collector - && self.gpu_experience_collector.is_none() - { - if let candle_core::Device::Cuda(ref cuda_dev) = self.device { - use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector; - let stream = cuda_dev.cuda_stream(); - - // Read-lock agent to extract dueling network weights for GPU collector. - // Works with both Standard and RegimeConditional agents — for regime-conditional, - // we use the primary (trending) head's weights since the GPU kernel runs a single - // Q-network for fast experience collection. - let agent = self.agent.read().await; - let dqn_ref: Option<&crate::dqn::DQN> = match &*agent { - DQNAgentType::Standard(ref dqn) => Some(dqn), - DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()), - }; - let init_result = if let Some(dqn) = dqn_ref { - // Curiosity is optional — pass None when curiosity_weight is 0 / module absent. - // The collector uses zero-filled weights + curiosity_scale=0.0 in that case. - let curiosity_vars = self.curiosity_module.as_ref() - .map(|c| c.forward_model_vars()); - - // kernel_dims: (state_dim, market_dim, num_atoms_max) - // state_dim = network input_dim (e.g. 56 with OFI alignment) - // market_dim = raw feature buffer width (42 features: 40 base + 2 regime) - // num_atoms_max = max atom count for C51 distributional - let state_dim = agent.get_state_dim(); - let market_dim: usize = 42; // market features incl. ADX + CUSUM regime indicators - let num_atoms_max = (self.hyperparams.num_atoms as usize).max(51); - let kernel_dims = (state_dim, market_dim, num_atoms_max); - - let use_branching = self.hyperparams.use_branching; - - // Priority: branching > plain dueling > hybrid (distributional+dueling) - if let (true, Some(online_br), Some(target_br)) = ( - use_branching, - dqn.branching_q_network.as_ref(), - dqn.branching_target_network.as_ref(), - ) { - // Branching DQN: use branching network VarMaps. - // The collector's extract_dueling_weights_branching maps - // branch_0_fc → advantage slot; branches 1+2 go to BranchingWeightSet. - let cfg = online_br.config(); - let dims = ( - *cfg.shared_hidden_dims.first().unwrap_or(&256), - *cfg.shared_hidden_dims.get(1).unwrap_or(&256), - cfg.value_hidden_dim, - cfg.branch_hidden_dim, - ); - let alloc_episodes = { - use ml_core::memory_optimization::detect_gpu_hardware; - let configured = self.hyperparams.gpu_n_episodes; - if configured >= 128 { - match detect_gpu_hardware() { - Ok(hw) => configured.max(hw.optimal_n_episodes( - state_dim, - self.hyperparams.gpu_timesteps_per_episode, - )).min(0x8000), - Err(_) => configured, - } - } else { - configured - } - }; - Some(GpuExperienceCollector::new( - stream, - online_br.vars(), - target_br.vars(), - curiosity_vars, - self.hyperparams.initial_capital as f32, - self.hyperparams.avg_spread as f32, - self.hyperparams.cash_reserve_percent as f32, - dims, - kernel_dims, - alloc_episodes, - self.hyperparams.gpu_timesteps_per_episode, - true, // use_branching - )) - } else if let (Some(online), Some(target)) = ( - dqn.dueling_q_network.as_ref(), - dqn.dueling_target_network.as_ref(), - ) { - let cfg = online.config(); - let dims = ( - *cfg.shared_hidden_dims.first().unwrap_or(&256), - *cfg.shared_hidden_dims.get(1).unwrap_or(&256), - cfg.value_hidden_dim, - cfg.advantage_hidden_dim, - ); - // Dynamic episode count for buffer allocation. - // Only auto-scale when configured >= 128 (production). - // Smaller values indicate an explicit test override — respect them. - let alloc_episodes = { - use ml_core::memory_optimization::detect_gpu_hardware; - let configured = self.hyperparams.gpu_n_episodes; - if configured >= 128 { - match detect_gpu_hardware() { - // Cap at GPU experience collector's buffer limit (32768) - Ok(hw) => configured.max(hw.optimal_n_episodes( - state_dim, - self.hyperparams.gpu_timesteps_per_episode, - )).min(0x8000), - Err(_) => configured, - } - } else { - configured - } - }; - Some(GpuExperienceCollector::new( - stream, - online.vars(), - target.vars(), - curiosity_vars, - self.hyperparams.initial_capital as f32, - self.hyperparams.avg_spread as f32, - self.hyperparams.cash_reserve_percent as f32, - dims, - kernel_dims, - alloc_episodes, - self.hyperparams.gpu_timesteps_per_episode, - self.hyperparams.use_branching, - )) - } else if let (Some(online), Some(target)) = ( - dqn.dist_dueling_q_network.as_ref(), - dqn.dist_dueling_target_network.as_ref(), - ) { - let cfg = online.config(); - let dims = ( - *cfg.shared_hidden_dims.first().unwrap_or(&256), - *cfg.shared_hidden_dims.get(1).unwrap_or(&256), - cfg.value_hidden_dim, - cfg.advantage_hidden_dim, - ); - // Dynamic episode count for buffer allocation. - // Only auto-scale when configured >= 128 (production). - // Smaller values indicate an explicit test override — respect them. - let alloc_episodes = { - use ml_core::memory_optimization::detect_gpu_hardware; - let configured = self.hyperparams.gpu_n_episodes; - if configured >= 128 { - match detect_gpu_hardware() { - // Cap at GPU experience collector's buffer limit (32768) - Ok(hw) => configured.max(hw.optimal_n_episodes( - state_dim, - self.hyperparams.gpu_timesteps_per_episode, - )).min(0x8000), - Err(_) => configured, - } - } else { - configured - } - }; - Some(GpuExperienceCollector::new( - stream, - online.vars(), - target.vars(), - curiosity_vars, - self.hyperparams.initial_capital as f32, - self.hyperparams.avg_spread as f32, - self.hyperparams.cash_reserve_percent as f32, - dims, - kernel_dims, - alloc_episodes, - self.hyperparams.gpu_timesteps_per_episode, - self.hyperparams.use_branching, - )) - } else { - return Err(anyhow::anyhow!( - "GPU experience collector FAILED: no dueling, hybrid, or branching networks — \ - DQN agent has no Q-network. This is a configuration bug." - )); - } - } else { - None - }; - drop(agent); - - if let Some(result) = init_result { - match result { - Ok(mut collector) => { - // Lazy-init GPU monitoring reducer on same stream as collector - if self.gpu_monitoring.is_none() { - match crate::cuda_pipeline::gpu_monitoring::GpuMonitoringReducer::new(collector.stream()) { - Ok(mon) => { - info!("GPU monitoring reducer initialized"); - self.gpu_monitoring = Some(mon); - } - Err(e) => { - return Err(anyhow::anyhow!( - "GPU monitoring reducer init FAILED: {e} — metrics pipeline must be GPU-resident" - )); - } - } - } - // Upload OFI features to GPU if available. - // This allows the kernel to populate state[45..53] with real - // OFI data instead of zeros, matching the CPU training path. - if let Some(ref ofi) = self.ofi_features { - let flat: Vec = ofi.iter() - .flat_map(|f| f.iter().map(|&v| v as f32)) - .collect(); - collector.upload_ofi_features(&flat)?; - } - self.gpu_experience_collector = Some(collector); - } - Err(e) => { - return Err(anyhow::anyhow!( - "GPU experience collector init FAILED (no CPU fallback): {e}" - )); - } - } - } - } - } - - let mut epoch_loss = 0.0; - let mut epoch_q_value = 0.0; - let mut epoch_gradient_norm = 0.0; - - // **WAVE 3 FIX #2: Two-Phase Feature Normalization** (Enhanced with configurable ratio) - // - // Phase 1 (epochs 0-N): Collect feature statistics - // - Build mean/std using Welford's algorithm (numerically stable) - // - N = min(epochs * ratio, max_epochs) - // - Default: min(epochs * 0.3, 10) → 10 epochs for <34 total epochs - // - No normalization applied yet - // - // Phase 2 (epochs N+1 onwards): Apply z-score normalization - // - Normalize all 42 market features to mean=0, std=1 - // - Portfolio features (indices 42-44) added separately via PortfolioTracker - // - Expected impact: Q-values reduced from ±10,000 to ±375 (27x improvement) - - // Calculate stats collection epochs using flexible formula - let _stats_collection_epochs = { - let ratio_based = (self.hyperparams.epochs as f32 * self.hyperparams.feature_stats_collection_ratio) as usize; - let capped = match self.hyperparams.max_feature_stats_epochs { - Some(max_epochs) => ratio_based.min(max_epochs), - None => ratio_based, // No cap if None - }; - capped.max(1) // Always collect at least 1 epoch - }; - - // Phase 3: GPU experience collection via zero-roundtrip CUDA kernel - // Collects N episodes × L timesteps entirely on GPU, then bulk-inserts into replay buffer. - // CUDA builds: no fallback — failure is a hard error. - // Branching DQN: the fused kernel now supports branching via use_branching flag - // + 16 extra weight pointers for order/urgency heads. - #[cfg(feature = "cuda")] - let gpu_experiences_collected = if let ( - Some(ref mut collector), - Some(ref features_buf), - Some(ref targets_buf), - ) = ( - &mut self.gpu_experience_collector, - &self.features_raw_cuda, - &self.targets_raw_cuda, - ) { - use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig; - - // Dynamic episode count: use GPU hardware to maximize SM utilization. - // Falls back to configured value when GPU detection fails. - // Only auto-scale when configured >= 128 (production). - // Smaller values indicate an explicit test override — respect them. - let raw_sd = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; - let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &self.device); - // Cache n_episodes on first epoch — detect_gpu_hardware() forks - // nvidia-smi (~5-10ms), and SM count / device name never change. - let n_episodes = if let Some(cached) = self.cached_n_episodes { - cached - } else { - use ml_core::memory_optimization::detect_gpu_hardware; - let configured = self.hyperparams.gpu_n_episodes; - let computed = if configured >= 128 { - match detect_gpu_hardware() { - Ok(hw) => { - let optimal = hw.optimal_n_episodes( - aligned_sd, - self.hyperparams.gpu_timesteps_per_episode, - ); - let chosen = configured.max(optimal).min(4096); - if chosen != configured { - info!( - "GPU auto-scaled n_episodes: {} → {} (SMs={}, VRAM={:.0}MB)", - configured, chosen, hw.sm_count, hw.free_memory_mb - ); - } - chosen as i32 - } - Err(_) => configured as i32, - } - } else { - configured as i32 - }; - self.cached_n_episodes = Some(computed); - computed - }; - let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32; - let total_bars = training_data.len() as i32; - let usable_bars = (total_bars - timesteps).max(1); - let stride = (usable_bars / n_episodes).max(1); - let episode_starts: Vec = (0..n_episodes) - .map(|i| (i * stride).rem_euclid(usable_bars)) - .collect(); - - // Reset per-episode state before each epoch - if let Err(e) = collector.reset_episodes( - self.hyperparams.initial_capital as f32, - self.hyperparams.avg_spread as f32, - self.hyperparams.cash_reserve_percent as f32, - ) { - return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}")); - } else { - let agent = self.agent.read().await; - let epsilon = agent.get_epsilon(); - drop(agent); - - let config = ExperienceCollectorConfig { - n_episodes, - timesteps_per_episode: timesteps, - total_bars, - epsilon, - gamma: self.hyperparams.gamma as f32, - // Position limit for portfolio simulation + action masking - max_position: self.max_position as f32, - // Action masking: filter invalid exposure actions in GPU kernel - enable_action_masking: self.enable_action_masking, - // Zero out curiosity bonus when curiosity module is disabled - curiosity_scale: if self.curiosity_module.is_some() { 1.0 } else { 0.0 }, - hold_reward: self.hyperparams.hold_penalty.abs() as f32, - tx_cost_multiplier: self.hyperparams.transaction_cost_multiplier as f32, - // D1: UCB count-bonus — same coefficient as CPU path - count_bonus_coefficient: self.hyperparams.count_bonus_coefficient - .unwrap_or(0.0) as f32, - // D2: Q-value clipping — conservative bounds (GPU numeric stability) - q_clip_min: -500.0, - q_clip_max: 500.0, - // D3: Huber TD — use configured delta for robust PER priorities - huber_kappa: if self.hyperparams.use_huber_loss { - self.hyperparams.huber_delta as f32 - } else { - 0.0 // Disabled: fall back to raw L1 - }, - // D5: NoisyNet — factorized noise exploration in GPU kernel - use_noisy_nets: self.hyperparams.use_noisy_nets, - noisy_sigma_init: self.hyperparams.noisy_sigma_init as f32, - // D6: C51 distributional — correct atom-distribution forward pass - use_distributional: self.hyperparams.use_distributional, - num_atoms: self.hyperparams.num_atoms as i32, - v_min: self.hyperparams.v_min as f32, - v_max: self.hyperparams.v_max as f32, - // Fill simulation: order routing + stochastic fill check - fill_median_spread: self.hyperparams.avg_spread as f32, - fill_median_vol: self.median_vol as f32, - fill_ioc_fill_prob: 0.85, - fill_limit_fill_min: 0.30, - fill_limit_fill_max: 0.80, - fill_spread_cost_frac: 0.50, - fill_spread_capture_frac: 0.50, - fill_simulation_enabled: self.median_vol > 0.0, - // DSR + N-step: GPU-side reward shaping - use_dsr: self.hyperparams.use_dsr, - dsr_eta: self.hyperparams.dsr_eta as f32, - n_steps: self.hyperparams.n_steps as i32, - ..Default::default() - }; - - // Check GPU PER status before kernel launch to choose the - // zero-roundtrip path (DtoD) vs the CPU download path. - let use_gpu_per = { - let agent = self.agent.read().await; - agent.memory().is_gpu_prioritized() - }; - - if use_gpu_per { - // ---- Zero-roundtrip GPU path ---- - // States and rewards stay on GPU via cuMemcpyDtoDAsync. - // Monitoring deferred to GpuMonitoringReducer (no per-launch downloads). - match collector.collect_experiences_gpu( - features_buf, targets_buf, &episode_starts, &config, &self.device, - ) { - Ok(gpu_batch) => { - let count = gpu_batch.n_episodes * gpu_batch.timesteps; - info!("GPU collected {} experiences (zero-roundtrip, {} episodes × {} timesteps)", - count, gpu_batch.n_episodes, gpu_batch.timesteps); - - // Deferred monitoring: reduce on GPU, download at epoch end - if let Some(ref mut mon) = self.gpu_monitoring { - if let Err(e) = mon.reduce(collector.rewards_gpu(), collector.actions_gpu(), count) { - debug!("GPU monitoring reduce failed (non-fatal): {e}"); - } - } - - // Train curiosity forward model on GPU-resident experience data - // (zero CPU traffic — reads directly from collector buffers) - if count > 0 { - if let Err(e) = collector.train_curiosity_gpu( - gpu_batch.n_episodes, gpu_batch.timesteps, - ) { - debug!("GPU curiosity training failed (non-fatal): {e}"); - } - } - - if count > 0 { - let agent = self.agent.read().await; - agent.insert_batch_tensors( - &gpu_batch.states, - &gpu_batch.next_states, - &gpu_batch.actions, - &gpu_batch.rewards, - &gpu_batch.dones, - ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; - } - true - } - Err(e) => { - return Err(anyhow::anyhow!( - "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" - )); - } - } - } else { - // ---- GPU download path (non-zero-roundtrip, used when PER is CPU-based) ---- - match collector.collect_experiences(features_buf, targets_buf, &episode_starts, &config) { - Ok(batch) => { - for &reward in &batch.rewards { - self.pnl_history.push_back(reward as f64); - if self.pnl_history.len() > 1000 { - self.pnl_history.pop_front(); - } - monitor.track_reward(reward); - } - for &action_idx in &batch.actions { - let exp_idx = action_idx.clamp(0, 4) as usize; - // GPU kernel only selects 5 exposure actions; route to - // factored (45-action) via heuristic so monitoring tracks - // the composite index consistently with the CPU path. - // track_action() updates both action_counts AND factored_action_counts. - if let Ok(exp_level) = crate::dqn::action_space::ExposureLevel::from_index(exp_idx) { - let factored = self.route_action(exp_level, self.hyperparams.avg_spread as f32); - monitor.track_action(&factored); - } else { - monitor.track_action_by_exposure(exp_idx); - } - } - - let raw_dim = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; - let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - let count = batch.states.len() / aligned_dim; - info!("GPU collected {} experiences ({} episodes × {} timesteps)", - count, batch.n_episodes, batch.timesteps); - - // Non-GPU-PER path: download + convert to Experience structs - let experiences = gpu_batch_to_experiences(&batch, aligned_dim); - self.store_experiences_batch(experiences).await?; - true - } - Err(e) => { - return Err(anyhow::anyhow!( - "GPU experience collection FAILED (no CPU fallback): {e}" - )); - } - } - } - } - } else { - false - }; - - #[cfg(not(feature = "cuda"))] - let gpu_experiences_collected = false; - - // Free DqnGpuData when GPU experience collector is active — it's only needed - // for the non-CUDA experience loop. On a 4 GB card this reclaims ~90 MB of VRAM - // (BF16 features + F32 targets) plus the candle-retained F32→BF16 conversion - // block (~150 MB). The raw cudarc buffers (features_raw_cuda / targets_raw_cuda) - // remain for the GPU experience collector. - if gpu_experiences_collected && self.gpu_data.is_some() { - info!("GPU experience collector active — releasing DqnGpuData to reclaim VRAM"); - self.gpu_data = None; - // Also release the buffer pool's CPU staging vecs (up to 40 MB RSS) - if let Some(pool) = self.buffer_pool.take() { - drop(pool); - } - // NOTE: features_raw_cuda / targets_raw_cuda must NOT be released here. - // The GPU experience collector kernel reads them on EVERY epoch. - // Releasing them after epoch 1 causes a hard error on epoch 2+ - // because the destructuring at the top of the experience collection - // block requires both buffers to be Some. - } - - // CUDA builds: GPU experience collector is MANDATORY — no CPU fallback exists. - // The entire CPU experience loop is compiled out (#[cfg(not(feature = "cuda"))]). - #[cfg(feature = "cuda")] - if !gpu_experiences_collected { - return Err(anyhow::anyhow!( - "GPU experience collector MUST be active for CUDA training. \ - No CPU fallback path exists in CUDA builds. \ - Set enable_gpu_experience_collector=true \ - or check GPU collector initialization errors above." - )); - } - - // CPU-only experience collection: compiled out entirely on CUDA builds. - // On non-CUDA builds, this is the only experience collection path. - #[cfg(not(feature = "cuda"))] - if !gpu_experiences_collected { - - // **PHASE 1: GPU-Optimized Experience Collection with Batched Action Selection** - // Fill replay buffer with batched action selection (125× fewer GPU kernel launches) - const ACTION_BATCH_SIZE: usize = 128; - let total_samples = training_data.len(); - let num_batches = (total_samples + ACTION_BATCH_SIZE - 1) / ACTION_BATCH_SIZE; - - for batch_idx in 0..num_batches { - let batch_start = batch_idx * ACTION_BATCH_SIZE; - let batch_end = ((batch_idx + 1) * ACTION_BATCH_SIZE).min(total_samples); - let batch_indices: Vec = (batch_start..batch_end).collect(); - - // Pre-allocate batch collectors for deferred lock acquisition - // Instead of acquiring agent write/read lock per sample, we collect - // actions and experiences and batch-write after the inner loop. - let batch_len = batch_end - batch_start; - let mut batch_actions_to_track: Vec = Vec::with_capacity(batch_len); - let mut batch_experiences_to_store: Vec = Vec::with_capacity(batch_len); - - // Build batch states — GPU path skips ~770 Vec allocs per batch - #[cfg(feature = "cuda")] - let (batch_tensor, states) = if let Some(ref gpu_data) = self.gpu_data { - // GPU path: build [batch_size, state_dim] directly from pre-uploaded features - let current_price_f32 = { - let t = gpu_data.bar_target_values(batch_start)?; - if t[2] != 0.0 { t[2] } else { t[0] } - }; - let portfolio_features = self.portfolio_tracker.get_portfolio_features(current_price_f32); - let bt = gpu_data.build_batch_states( - batch_start, - batch_end - batch_start, - &portfolio_features, - &self.device, - )?; - - // Still need TradingStates for reward calculation and experience storage - let cpu_states: Result> = batch_indices - .iter() - .map(|&i| { - let target = &training_data[i].1; - let current_close = if target.len() >= 2 { - target[0] - } else { - training_data[i].0[3] - }; - let close_price = rust_decimal::Decimal::try_from(current_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - self.feature_vector_to_state_with_ofi(&training_data[i].0, Some(close_price), Some(i)) - }) - .collect(); - (Some(bt), cpu_states?) - } else { - // CPU fallback: original path - let cpu_states: Result> = batch_indices - .iter() - .map(|&i| { - let target = &training_data[i].1; - let current_close = if target.len() >= 2 { - target[0] - } else { - training_data[i].0[3] - }; - let close_price = rust_decimal::Decimal::try_from(current_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - self.feature_vector_to_state_with_ofi(&training_data[i].0, Some(close_price), Some(i)) - }) - .collect(); - (None, cpu_states?) - }; - #[cfg(not(feature = "cuda"))] - let states: Vec = { - let cpu_states: Result> = batch_indices - .iter() - .map(|&i| { - let target = &training_data[i].1; - let current_close = if target.len() >= 2 { - target[0] - } else { - training_data[i].0[3] - }; - let close_price = rust_decimal::Decimal::try_from(current_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - self.feature_vector_to_state_with_ofi(&training_data[i].0, Some(close_price), Some(i)) - }) - .collect(); - cpu_states? - }; - - // Batched action selection — GPU tensor path skips state→Vec→flatten→Tensor ✅ - // gpu_handled_fill: when true, GPU kernel already did routing + fill simulation. - // Inner loop skips CPU route_action() + simulate_fill() to avoid double-work. - #[cfg(feature = "cuda")] - let (actions, gpu_handled_fill) = if let Some(ref bt) = batch_tensor { - self.select_actions_batch_gpu(bt, batch_start).await? - } else { - (self.select_actions_batch(&states).await?, false) - }; - #[cfg(not(feature = "cuda"))] - let (actions, gpu_handled_fill) = (self.select_actions_batch(&states).await?, false); - - // GPU portfolio simulation for the batch (if CUDA available). - // When gpu_handled_fill=true, actions are post-fill so GPU sim PnL is correct. - // When gpu_handled_fill=false, actions are pre-fill — GPU sim would simulate - // unfilled trades, so we fall back to CPU portfolio tracking. - #[cfg(feature = "cuda")] - let gpu_sim_results = if gpu_handled_fill { - if let (Some(ref mut sim), Some(ref targets_buf)) = - (&mut self.gpu_portfolio_sim, &self.targets_raw_cuda) - { - let action_indices: Vec = actions.iter().map(|a| a.to_index() as i32).collect(); - match sim.simulate_batch(targets_buf, &action_indices, batch_start) { - Ok(result) => Some(result), - Err(e) => { - return Err(anyhow::anyhow!("GPU portfolio sim FAILED (no CPU fallback): {e}")); - } - } - } else { - None - } - } else { - None - }; - - // Store experiences with batched actions - for (idx_in_batch, &i) in batch_indices.iter().enumerate() { - let state = &states[idx_in_batch]; - let raw_action = actions[idx_in_batch]; - - // Extract GPU sim data for this sample (reward, portfolio features, done flag). - // Available when gpu_handled_fill=true and GPU portfolio sim succeeded. - // When Some, skips CPU portfolio execution and RewardFunction::calculate_reward(). - #[cfg(feature = "cuda")] - let gpu_sim_sample: Option<(f32, [f32; 3], bool)> = - gpu_sim_results.as_ref().and_then(|r| { - Some(( - *r.rewards.get(idx_in_batch)?, - [ - *r.portfolio_features.get(idx_in_batch * 3)?, - *r.portfolio_features.get(idx_in_batch * 3 + 1)?, - *r.portfolio_features.get(idx_in_batch * 3 + 2)?, - ], - *r.done_flags.get(idx_in_batch)? != 0, - )) - }); - #[cfg(not(feature = "cuda"))] - let gpu_sim_sample: Option<(f32, [f32; 3], bool)> = None; - - // Phase C/C+: Order routing + fill simulation - // When gpu_handled_fill=true, the GPU fused kernel already did: - // 1. epsilon-greedy action selection - // 2. order routing (spread/vol → order type + urgency) - // 3. fill simulation (splitmix64 hash → fill/no-fill) - // 4. unfilled → Flat override - // Skip CPU routing + fill to avoid double-work. - let action = if gpu_handled_fill { - raw_action // Already routed + fill-checked by GPU kernel - } else { - let current_spread = { - let features = &training_data[i].0; - // HL spread estimate: ln(high/low) ≈ features[1] - features[2] - (features[1] - features[2]).abs() as f32 - }; - let routed_action = if self.hyperparams.use_branching { - raw_action // Branching: order/urgency learned by network - } else { - self.route_action(raw_action.exposure, current_spread) - }; - // Phase C: Fill simulation gate — check if order fills - let (fill_action, _fill_result) = self.simulate_fill(routed_action, i); - fill_action - }; - - // Update volatility EMA from close log return (features[3]) - let vol_sample = training_data[i].0[3].abs(); - if vol_sample > 0.0 && vol_sample < 1.0 { - // Fast EMA (α=0.01) for current vol, slow EMA (α=0.001) for median - self.vol_ema = 0.99 * self.vol_ema + 0.01 * vol_sample; - self.median_vol = 0.999 * self.median_vol + 0.001 * vol_sample; - } - - // Use pre-uploaded GPU targets if available, otherwise fall back to CPU - // WAVE 3 BUG FIX: Use raw prices (indices 2,3) for barrier tracker, preprocessed (indices 0,1) for rewards - let (current_close_raw, next_close_raw, current_close, next_close) = - if let Some(ref gpu_data) = self.gpu_data { - let t = gpu_data.bar_target_values(i)?; - let cc = t[0] as f64; - let nc = t[1] as f64; - let ccr = if t[2] != 0.0 { t[2] as f64 } else { cc }; - let ncr = if t[3] != 0.0 { t[3] as f64 } else { cc }; - (ccr, ncr, cc, nc) - } else { - let target = &training_data[i].1; - let current_close_raw = if target.len() >= 4 { - target[2] // Raw price for barrier tracker - } else if target.len() >= 2 { - target[0] // Fallback to preprocessed for old data - } else { - training_data[i].0[3] - }; - let next_close_raw = if target.len() >= 4 { - target[3] // Raw price for barrier tracker - } else if target.len() >= 2 { - target[1] // Fallback to preprocessed for old data - } else { - current_close_raw - }; - let current_close = if target.len() >= 2 { - target[0] // Preprocessed for reward calculation - } else { - training_data[i].0[3] - }; - let next_close = if target.len() >= 2 { - target[1] // Preprocessed for reward calculation - } else { - current_close - }; - (current_close_raw, next_close_raw, current_close, next_close) - }; - - // Get next state (BUG #42: made mutable to update portfolio features after trade execution) - let mut next_state = if i + 1 < training_data.len() { - let next_close_price = rust_decimal::Decimal::try_from(next_close) - .unwrap_or(rust_decimal::Decimal::ZERO); - self.feature_vector_to_state_with_ofi( - &training_data[i + 1].0, - Some(next_close_price), - Some(i + 1), - )? - } else { - state.clone() - }; - - // WAVE 3 AGENT 2: Simulated Position Tracking for Barrier Episodes - // Root Cause: BUG #8 prevents portfolio execution during experience collection, - // so positions never change and barriers never trigger (0% barrier exits). - // Solution: Map action exposure to simulated position for barrier tracking. - let simulated_position = match action.exposure { - crate::dqn::action_space::ExposureLevel::Long100 => 1.0, - crate::dqn::action_space::ExposureLevel::Long50 => 0.5, - crate::dqn::action_space::ExposureLevel::Flat => 0.0, - crate::dqn::action_space::ExposureLevel::Short50 => -0.5, - crate::dqn::action_space::ExposureLevel::Short100 => -1.0, - }; - - // Override next_state portfolio features with simulated position - // Note: This ONLY affects barrier tracking, NOT the actual state stored in experience - let mut next_state_with_sim_position = next_state.clone(); - if next_state_with_sim_position.portfolio_features.len() >= 2 { - next_state_with_sim_position.portfolio_features[1] = simulated_position; - } - - // WAVE 1.1 + WAVE 3: Triple Barrier Position Tracking with Simulated Positions - // WAVE P3 FIX: Use previous_simulated_position instead of state.portfolio_features[1] - // Bug: state.portfolio_features[1] is always 0.0 during experience collection (BUG #8 fix) - // This caused position_changed=true on EVERY step, creating new tracker each iteration - let current_position = self.previous_simulated_position; - // WAVE 3: Use simulated position from action intent (not portfolio tracker) - let next_position = next_state_with_sim_position.portfolio_features.get(1).unwrap_or(&0.0); - let position_changed = (next_position - current_position).abs() > 0.01; - - // Update previous position for next iteration (WAVE P3 FIX) - self.previous_simulated_position = simulated_position; - - // Start tracking if position changed and no active tracker - if position_changed && self.active_position_tracker.is_none() && next_position.abs() > 0.01 { - // Use conservative barrier configuration - let config = BarrierConfig::conservative(); - - // Convert price to cents (multiply by 100) - // WAVE 3 BUG FIX: Use raw price (not preprocessed) to avoid divide-by-zero - let entry_price_cents = (current_close_raw * 100.0) as u64; - - // Skip barrier tracking if price is zero (corrupt/missing data) - if entry_price_cents == 0 { - debug!("Skipping barrier tracking at step {}: entry price is zero", i); - } else { - - // Use step index as timestamp (nanoseconds) - // Each step = 1 second for simplicity - let entry_timestamp_ns = (i as u64) * 1_000_000_000; - - // Start tracking the position - match self.triple_barrier.write().await.start_tracking( - config, - entry_price_cents, - entry_timestamp_ns, - ) { - Ok(tracker_id) => { - self.active_position_tracker = Some(tracker_id); - debug!( - "WAVE 1.1: Started triple barrier tracking at step {}, price=${:.2}, position={:.2}", - i, current_close, next_position - ); - }, - Err(e) => { - warn!("WAVE 1.1: Failed to start triple barrier tracking: {}", e); - } - } - } // close else block (entry_price_cents != 0) - } - - // WAVE P2: Check for barrier exits on each step - // Changed to Option to distinguish "no barrier" (None) from "time expiry" (Some(0)) - let mut barrier_label: Option = None; // None = no barrier, Some(0/1/-1) = barrier hit - if let Some(tracker_id) = self.active_position_tracker { - // Create price point for current step - // WAVE 3 BUG FIX: Use raw price (not preprocessed) for barrier calculations - let price_cents = (next_close_raw * 100.0) as u64; - let timestamp_ns = ((i + 1) as u64) * 1_000_000_000; - let price_point = PricePoint::new(price_cents, timestamp_ns); - - // Check if any barrier was hit - if let Some(event_label) = self.triple_barrier.write().await.update_tracker( - tracker_id, - price_point, - ) { - // WAVE P2: Use Option to distinguish barrier events from no-barrier - barrier_label = Some(event_label.label_value); // Some(1/0/-1) indicates barrier hit - self.active_position_tracker = None; // Clear tracker on exit - - // WAVE P2: Reset portfolio on barrier exit (position closes) - self.portfolio_tracker.reset(); - self.reward_fn.reset_dsr(); - - debug!( - "WAVE P2: Barrier-driven episode end at step {}: {:?}, label={}, return_bps={}, portfolio reset", - i, event_label.barrier_result, barrier_label.unwrap_or(0), event_label.return_bps - ); - } - } - - // Portfolio execution + next_state update. - // GPU path: GPU sim already computed PnL, portfolio features, and done flags. - // CPU path: execute trade on portfolio tracker (BUG #42 FIX). - if let Some((_, ref gpu_features, _)) = gpu_sim_sample { - // GPU sim computed portfolio features for this sample. - // Skip CPU portfolio tracker execution — GPU kernel already did - // the full position management, mark-to-market, and PnL calculation. - if next_state.portfolio_features.len() >= 3 { - next_state.portfolio_features[0] = gpu_features[0]; // Normalized value - next_state.portfolio_features[1] = gpu_features[1]; // Normalized position - // portfolio_features[2] is spread, leave unchanged (already set) - } - } else { - // CPU portfolio execution (BUG #42 FIX) - let position_before = self.portfolio_tracker.current_position(); - let target_exposure = action.target_exposure() as f32; - - let current_price_f32 = current_close as f32; - let max_position_f32 = self.max_position as f32; - self.portfolio_tracker.execute_action(action, current_price_f32, max_position_f32); - - let position_after = self.portfolio_tracker.current_position(); - - if i < 100 { - let action_idx = action.exposure as usize; - debug!( - "SIGN_DIAG step={} action={} target_exp={:+.3} pos_before={:+.4} pos_after={:+.4}", - i, action_idx, target_exposure, position_before, position_after - ); - } - - let next_price_f32 = next_close as f32; - let updated_portfolio_features = self.portfolio_tracker.get_portfolio_features(next_price_f32); - - if next_state.portfolio_features.len() >= 3 { - next_state.portfolio_features[0] = updated_portfolio_features[0]; - next_state.portfolio_features[1] = updated_portfolio_features[1]; - } - - if i % 100 == 0 || i < 10 { - let current_value = state.portfolio_features.get(0).unwrap_or(&1.0); - let new_value = updated_portfolio_features[0]; - let new_position = updated_portfolio_features[1]; - let reward_signal = if *current_value > 0.0 { - ((new_value - current_value) / current_value) * 100.0 - } else { - 0.0 - }; - debug!( - "BUG42_FIX: step={}, action={:?}, current_value={:.6}, new_value={:.6}, position={:.4}, reward_signal={:.4}%", - i, action, current_value, new_value, new_position, reward_signal - ); - } - } - - // Track action for diversity penalty - self.recent_actions.push_back(action); - if self.recent_actions.len() > 100 { - self.recent_actions.pop_front(); - } - - // WAVE 1.2 SAFETY #5: Action Diversity Monitor (every 100 steps) - self.safety_step_counter += 1; - // Track by composite factored action index (0-44) - let factored_idx = action.exposure as usize * 9 - + action.order as usize * 3 - + action.urgency as usize; - *self.safety_action_counts.entry(factored_idx).or_insert(0) += 1; - - if self.safety_step_counter % 100 == 0 { - let total_actions: usize = self.safety_action_counts.values().sum(); - let unique_actions = self.safety_action_counts.len(); - let diversity = unique_actions as f32 / 45.0; // 45 factored actions - - if diversity < 0.5 { - let msg = format!( - "SAFETY: Action diversity below 50% at step {}: {:.1}% ({}/{} factored actions used in last {} steps)", - self.safety_step_counter, diversity * 100.0, unique_actions, 45, total_actions - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{}", msg); - }, - } - } - - // Reset counts every 1000 steps - if self.safety_step_counter % 1000 == 0 { - self.safety_action_counts.clear(); - } - } - - // WAVE 1.2 SAFETY #6: Memory Monitor (every 10 steps) - if self.safety_step_counter % 10 == 0 { - if let Device::Cuda(_) = &self.device { - let memory_manager = self.safety_memory_manager.read().await; - let current_usage = memory_manager.get_memory_usage(&self.device); - // Use 4GB as limit for RTX 3050 Ti (memory_manager.get_memory_limit is private) - let limit = 4_000_000_000_usize; // 4GB in bytes - let usage_ratio = current_usage as f64 / limit as f64; - - if usage_ratio > 0.9 { - let usage_gb = current_usage as f64 / 1e9; - let limit_gb = limit as f64 / 1e9; - debug!( - "SAFETY: GPU memory usage high at step {}: {:.1}% ({:.2} GB / {:.2} GB)", - self.safety_step_counter, - usage_ratio * 100.0, - usage_gb, - limit_gb - ); - } - } - } - - // Calculate reward: GPU path uses pre-computed PnL, CPU path uses RewardFunction. - // Both paths apply barrier scaling and risk adjustment on CPU (lightweight scalar ops). - let (raw_reward, reward) = if let Some((gpu_reward, _, _)) = gpu_sim_sample { - // GPU portfolio sim already computed PnL reward with risk penalty. - // Apply barrier scaling + DSR on top (CPU scalar ops). - let raw = gpu_reward as f64; - let barrier_scaled = if let Some(label) = barrier_label { - let reward_dec = rust_decimal::Decimal::try_from(raw) - .unwrap_or(rust_decimal::Decimal::ZERO); - let scaled = self.reward_fn.apply_triple_barrier_scaling(reward_dec, label); - scaled.to_f64().unwrap_or(0.0) - } else { - raw - }; - let risk_adjusted = if self.hyperparams.use_dsr { - barrier_scaled - } else { - self.calculate_risk_adjusted_reward(barrier_scaled) - }; - (raw, risk_adjusted as f32) - } else { - // CPU reward path: full RewardFunction pipeline - let recent_actions_vec: Vec = - self.recent_actions.iter().copied().collect(); - let reward_decimal = self.reward_fn.calculate_reward( - action, - state, - &next_state, - &recent_actions_vec, - )?; - let raw = reward_decimal.to_f64().unwrap_or(0.0); - - let barrier_scaled = if let Some(label) = barrier_label { - let scaled = self.reward_fn.apply_triple_barrier_scaling(reward_decimal, label); - scaled.to_f64().unwrap_or(0.0) - } else { - raw - }; - - let risk_adjusted = if self.hyperparams.use_dsr { - barrier_scaled - } else { - self.calculate_risk_adjusted_reward(barrier_scaled) - }; - - // Curiosity forward model training (C1: reward not used, only weight update) - if let Some(ref mut curiosity) = self.curiosity_module { - let state_vec = state.to_vector(); - let state_tensor = Tensor::from_vec( state_vec.clone(), - (1, state_vec.len()), - &self.device - ).map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast state tensor to training dtype: {}", e))?; - - let next_state_vec = next_state.to_vector(); - let next_state_tensor = Tensor::from_vec( next_state_vec, - (1, state_vec.len()), - &self.device - ).map_err(|e| anyhow::anyhow!("Failed to create next_state tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast next_state tensor to training dtype: {}", e))?; - - // calculate_curiosity_reward also trains the forward model internally - let _ = curiosity.calculate_curiosity_reward( - &state_tensor, - action, - &next_state_tensor, - ).map_err(|e| anyhow::anyhow!("Curiosity forward model training failed: {}", e))?; - } - - (raw, risk_adjusted as f32) - }; - - // Calculate price return for volatility tracking - let price_return = (next_close - current_close) / current_close; - - // BUG #17 FIX: Use raw_reward to update trackers (prevents feedback loop) - self.update_risk_trackers(raw_reward, price_return); - - // WAVE 16S: Log adaptive features periodically - if i % 100 == 0 && i > 0 { - let kelly_frac = self.get_kelly_fraction(); - debug!( - "Step {}: Kelly={:.4}, Sharpe history={}, Vol returns={}", - i, kelly_frac, self.pnl_history.len(), self.volatility_returns.len() - ); - } - - // Track reward and action for monitoring - monitor.track_reward(reward); - monitor.track_action(&action); - // We'll track Q-values during training steps - - // BUG #8 FIX: DO NOT execute portfolio actions during experience collection - // Experience collection is for SIMULATION only (building replay buffer) - // Portfolio actions should ONLY be executed during: - // - Evaluation phase (compute_validation_loss) - // - Backtesting (separate EvaluationEngine) - // - NOT during training experience collection - // - // Portfolio features are already populated via feature_vector_to_state() - // which extracts them from FeatureVector (Bug #2 fix is separate) - - // Collect action for batched tracking (deferred to end of batch) - batch_actions_to_track.push(action); - - // WAVE P2: Barrier-Based Episode Termination - // Episodes end on THREE conditions: - // 1. Barrier hit (profit target, stop loss, or time expiry) - // 2. Fixed episode length boundary (fallback for compatibility) - // 3. Data boundary reached - let barrier_done = barrier_label.is_some(); // Any barrier event (Some(1/0/-1)) - let gpu_sim_done = gpu_sim_sample.map(|(_, _, d)| d).unwrap_or(false); - let time_done = (i + 1) % EPISODE_LENGTH == 0; - let data_done = i + 1 >= training_data.len(); - let done = barrier_done || gpu_sim_done || time_done || data_done; - - // Log barrier-driven terminations for analysis - if barrier_done { - let label_value = barrier_label.unwrap_or(0); - debug!( - "WAVE P2: Episode ended via barrier at step {} ({}): label={}", - i + 1, - match label_value { - 1 => "Profit Target", - -1 => "Stop Loss", - 0 => "Time Expiry", - _ => "Unknown", - }, - label_value - ); - } - - // P1 FIX: Reset portfolio at episode boundaries (updated for barrier termination) - // Note: Portfolio already reset in barrier detection block - // GPU sim done: GPU kernel already reset its internal portfolio state. - // This handles time/data boundary resets on the CPU tracker. - if done && !barrier_done && !gpu_sim_done { - self.portfolio_tracker.reset(); - self.reward_fn.reset_dsr(); - - let episode_num = ((i + 1) / EPISODE_LENGTH) + 1; - let total_episodes = (training_data.len() + EPISODE_LENGTH - 1) / EPISODE_LENGTH; - debug!( - "Episode boundary at sample {}/{} (time={}, data={}), portfolio reset (episode {}/{})", - i + 1, training_data.len(), time_done, data_done, episode_num, total_episodes - ); - } - - // WAVE P2: Track episode end for statistics - if done { - monitor.track_episode_end(i, barrier_label); - } - - // Build experience (with optional n-step accumulation) - // Branching DQN: store factored index (0-44) for per-branch decomposition - // Standard DQN: store exposure index (0-4) only - let action_idx = if self.hyperparams.use_branching { - action.to_index() as u8 - } else { - action.exposure as u8 - }; - let experience = Experience::new( - state.to_vector(), - action_idx, - reward, - next_state.to_vector(), - done, - ); - - // WAVE 44: Multi-step returns integration - // Collect experiences into batch_experiences_to_store for deferred batch write - if self.nstep_buffer.is_some() { - // Extract buffer to avoid borrow checker issues - let Some(mut nstep_buf) = self.nstep_buffer.take() else { continue; }; - - // Collect n-step experience if buffer is full - if let Some(nstep_exp) = nstep_buf.add(experience) { - batch_experiences_to_store.push(nstep_exp); - } - - // Flush remaining experiences at episode end - if done { - batch_experiences_to_store.extend(nstep_buf.flush()); - nstep_buf.clear(); // Reset for next episode - } - - // Put buffer back - self.nstep_buffer = Some(nstep_buf); - } else { - // n_steps=1: Standard single-step experience - batch_experiences_to_store.push(experience); - } - } - - // Batch write: track all actions with a single write lock acquisition - // (reduces async RwLock contention from O(batch_size) to O(1) per batch) - { - let mut agent = self.agent.write().await; - agent.track_actions_batch(&batch_actions_to_track); - } - - // Batch write: store all experiences with a single read lock + inner mutex - self.store_experiences_batch(batch_experiences_to_store).await?; - } - - } // end #[cfg(not(feature = "cuda"))] if !gpu_experiences_collected - - // **PHASE 2: Batched Training from Replay Buffer** - // Now that buffer is populated, perform batched training - // This reduces train_step() calls from 1000×/epoch to ~8×/epoch (125× reduction) - - let batch_size = self.hyperparams.batch_size; - let num_training_steps = if self.can_train().await? { - // Calculate number of training steps based on dataset size and batch size - let full_steps = (training_data.len() / batch_size).max(1); - // Cap to max_training_steps_per_epoch if set (CI/smoke: 50-100 for fast validation) - let cap = self.hyperparams.max_training_steps_per_epoch; - if cap > 0 { full_steps.min(cap) } else { full_steps } - } else { - // Buffer not ready yet (early epochs) - 0 - }; - - let mut train_step_count = 0; - - // Check if GPU PER is active — determines zero-sync vs CPU training path - let is_gpu_per = { - let agent = self.agent.read().await; - agent.memory().is_gpu_prioritized() - }; - - if is_gpu_per { - // ═══ GPU-accumulation path: guard kernel accumulates loss/grad ═══ - // Fused check_and_accumulate() does NaN/loss-clip/grad-collapse - // checks AND accumulates loss_sum/grad_sum/step_count in device - // memory. Single mapped-memory readback at epoch boundary. - - // Lazy-init + reset guard accumulators for this epoch - #[cfg(feature = "cuda")] - { - if self.training_guard.is_none() && self.device.is_cuda() { - match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) { - Ok(g) => { - info!("GPU training guard initialized (epoch loop)"); - self.training_guard = Some(g); - } - Err(e) => return Err(anyhow::anyhow!("GPU training guard init: {e}")), - } - } - if let Some(ref mut guard) = self.training_guard { - guard.reset_accumulators() - .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?; - } - } - #[cfg(feature = "cuda")] - let guard_collapse_thresh = - self.hyperparams.learning_rate as f32 - * self.hyperparams.gradient_collapse_multiplier as f32; - #[cfg(feature = "cuda")] - let guard_past_warmup = { - // Use original buffer_size (before AutoReplaySizer) for warmup guard - let ws = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; - self.gradient_logging_step as u64 > ws - }; - - // Batch pre-sampling: sample K batches under one READ lock, train all - // under one WRITE lock. Reduces async lock acquisitions from 2×N to - // 2×ceil(N/K). Priority staleness across K steps is negligible (K=8). - const PREFETCH_K: usize = 32; - for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) { - let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps); - - // Pre-sample chunk of batches (single READ lock) - let batches = { - let agent = self.agent.read().await; - let buffer = agent.memory(); - let mut b = Vec::with_capacity(chunk_end - chunk_start); - for _ in chunk_start..chunk_end { - b.push(buffer.can_sample(self.current_batch_size).then(|| { - buffer - .sample(self.current_batch_size) - .map_err(|e| anyhow::anyhow!("Pre-sample: {e}")) - }).transpose()?); - } - b - }; - - // GPU train steps (single WRITE lock) — zero CPU sync - { - let mut agent = self.agent.write().await; - let accum_steps = self.hyperparams.gradient_accumulation_steps; - - if accum_steps <= 1 { - // Standard path: one forward+backward+optimizer per batch - for explicit_batch in batches { - let _gpu_result = match agent.train_step(explicit_batch) { - Ok(r) => r, - Err(e) => { - let msg = e.to_string(); - if msg.contains("Early stopping") || msg.contains("Gradient collapse") { - return Err(anyhow::anyhow!("{}", msg)); - } - return Err(anyhow::anyhow!("GPU train step FAILED (no CPU fallback): {e}")); - } - }; - - // Guard kernel: NaN/clip/collapse check + device-memory accumulation - #[cfg(feature = "cuda")] - if let Some(ref mut guard) = self.training_guard { - let gr = guard.check_and_accumulate( - &_gpu_result.loss_gpu, - &_gpu_result.grad_norm_gpu, - 1e6_f32, - guard_collapse_thresh, - !guard_past_warmup, - ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?; - if gr.halt_nan { - return Err(anyhow::anyhow!( - "NaN/Inf at step {}: loss={}, grad={}", - train_step_count, gr.raw_loss, gr.raw_grad_norm - )); - } - if gr.halt_grad_collapse { - agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - } - train_step_count += 1; - self.gradient_logging_step += 1; - } - } else { - // Gradient accumulation: group every accum_steps batches into - // one optimizer step. Effective batch = batch_size × accum_steps. - // Zero extra VRAM — same batch_size memory, more forward passes. - let mut batch_iter = batches.into_iter().peekable(); - - while batch_iter.peek().is_some() { - let mut accumulated_grads: Option = None; - #[cfg(feature = "cuda")] - let mut group_td_gpu: Vec = Vec::new(); - #[cfg(feature = "cuda")] - let mut group_idx_gpu: Vec = Vec::new(); - let mut accum_count: usize = 0; - - // Accumulate forward+backward over accum_steps mini-batches - for _ in 0..accum_steps { - let batch = match batch_iter.next() { - Some(b) => b, - None => break, - }; - - let result = match agent.compute_gradients(batch) { - Ok(r) => r, - Err(e) => { - let msg = e.to_string(); - if msg.contains("Early stopping") || msg.contains("Gradient collapse") { - return Err(anyhow::anyhow!("{}", msg)); - } - return Err(anyhow::anyhow!("GPU gradient compute FAILED: {e}")); - } - }; - - // Accumulate GradStore - let vars = agent.optimizer_vars() - .map_err(|e| anyhow::anyhow!("optimizer vars: {e}"))?; - crate::gradient_accumulation::accumulate_grads( - &mut accumulated_grads, - result.grads, - &vars, - ).map_err(|e| anyhow::anyhow!("grad accum: {e}"))?; - - // Guard kernel: NaN/clip check + device-memory accumulation - #[cfg(feature = "cuda")] - { - if let (Some(ref loss_t), Some(ref gn_t)) = - (&result.loss_tensor_gpu, &result.grad_norm_gpu) - { - if let Some(ref mut guard) = self.training_guard { - let gr = guard.check_and_accumulate( - loss_t, gn_t, 1e6_f32, - guard_collapse_thresh, !guard_past_warmup, - ).map_err(|e| anyhow::anyhow!("guard accum step: {e}"))?; - if gr.halt_nan { - return Err(anyhow::anyhow!( - "NaN/Inf in accum sub-step: loss={}, grad={}", - gr.raw_loss, gr.raw_grad_norm - )); - } - if gr.halt_grad_collapse { - agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - } - } - if let Some(td) = result.td_errors_gpu { - group_td_gpu.push(td); - } - if let Some(idx) = result.indices_gpu { - group_idx_gpu.push(idx); - } - } - - accum_count += 1; - } - - if accum_count == 0 { break; } - - // Scale and apply (single optimizer step) - if let Some(ref mut grads) = accumulated_grads { - let vars = agent.optimizer_vars() - .map_err(|e| anyhow::anyhow!("optimizer vars: {e}"))?; - crate::gradient_accumulation::scale_grads( - grads, &vars, 1.0 / accum_count as f64, - ).map_err(|e| anyhow::anyhow!("grad scale: {e}"))?; - - agent.apply_accumulated_gradients(grads) - .map_err(|e| anyhow::anyhow!("apply grads: {e}"))?; - } - - // Update priorities for all batches in this group - #[cfg(feature = "cuda")] - if !group_td_gpu.is_empty() && !group_idx_gpu.is_empty() { - let td_cat = Tensor::cat(&group_td_gpu, 0) - .map_err(|e| anyhow::anyhow!("TD cat: {e}"))?; - let idx_cat = Tensor::cat(&group_idx_gpu, 0) - .map_err(|e| anyhow::anyhow!("idx cat: {e}"))?; - agent.update_priorities_gpu(&idx_cat, &td_cat) - .map_err(|e| anyhow::anyhow!("PER update: {e}"))?; - } - agent.step_replay_buffer(); - - train_step_count += accum_count; - self.gradient_logging_step += accum_count; - } - } - } - } - - // ═══ Single epoch-boundary readback via guard accumulators ═══ - if train_step_count > 0 { - let n = train_step_count as f64; - - #[cfg(feature = "cuda")] - let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard { - let (al, ag) = guard.read_accumulators() - .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?; - (al as f32, ag as f32) - } else { - (0.0_f32, 0.0_f32) - }; - #[cfg(not(feature = "cuda"))] - let (avg_loss, avg_grad) = (0.0_f32, 0.0_f32); - - // Q-value estimation at epoch boundary — pipeline already synced, - // so this to_scalar is free (no additional GPU flush). - let avg_q = { - let mut agent = self.agent.write().await; - self.estimate_avg_q_value_with_early_stopping(&mut agent).await? - }; - - // Epoch-level safety checks (all per-step checks are redundant — - // NaN detected on GPU every 500 steps, grad already clipped) - if !avg_loss.is_finite() || !avg_grad.is_finite() { - let msg = format!( - "SAFETY: NaN/Inf in epoch {} avg — loss={:.6}, grad={:.6}", - epoch, avg_loss, avg_grad - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{} (continuing training)", msg); - }, - } - } - - // Loss history (epoch-level average instead of per-step) - self.safety_loss_history.push_back(avg_loss); - if self.safety_loss_history.len() > 30 { - self.safety_loss_history.pop_front(); - } - - epoch_loss = avg_loss as f64 * n; - epoch_q_value = avg_q * n; - epoch_gradient_norm = avg_grad as f64 * n; - - // Metrics aggregator - let current_avg_reward = if !monitor.reward_history.is_empty() { - monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 - } else { - 0.0 - }; - self.metrics_aggregator.record(avg_loss, avg_q as f32, current_avg_reward); - self.metrics_aggregator.record_gradient(avg_grad); - - if self.metrics_aggregator.should_log(&self.logging_config) { - let aggregated = self.metrics_aggregator.aggregate_and_clear(); - log_training_progress(&aggregated); - } - - // Gradient diagnostics at epoch boundary - { - let mut agent = self.agent.write().await; - agent.log_diagnostics(avg_grad) - .map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - - monitor.track_q_value_range(avg_q); - } - } else { - // ═══ Non-GPU-PER path: guard on CUDA ═══ - - // CUDA: use guard kernel's device-memory accumulator - #[cfg(feature = "cuda")] - { - if self.training_guard.is_none() && self.device.is_cuda() { - match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) { - Ok(g) => { - info!("GPU training guard initialized (non-PER path)"); - self.training_guard = Some(g); - } - Err(e) => return Err(anyhow::anyhow!("GPU training guard init: {e}")), - } - } - if let Some(ref mut guard) = self.training_guard { - guard.reset_accumulators() - .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?; - } - } - #[cfg(feature = "cuda")] - let guard_collapse_thresh2 = - self.hyperparams.learning_rate as f32 - * self.hyperparams.gradient_collapse_multiplier as f32; - #[cfg(feature = "cuda")] - let guard_past_warmup2 = { - // Use original buffer_size (before AutoReplaySizer) for warmup guard - let ws = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; - self.gradient_logging_step as u64 > ws - }; - - - // Batch pre-sampling: same K=8 chunking as GPU-PER path (see above) - const PREFETCH_K: usize = 32; - for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) { - let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps); - - let batches = { - let agent = self.agent.read().await; - let buffer = agent.memory(); - let mut b = Vec::with_capacity(chunk_end - chunk_start); - for _ in chunk_start..chunk_end { - b.push(buffer.can_sample(self.current_batch_size).then(|| { - buffer - .sample(self.current_batch_size) - .map_err(|e| anyhow::anyhow!("Pre-sample: {e}")) - }).transpose()?); - } - b - }; - - { - let mut agent = self.agent.write().await; - for explicit_batch in batches { - let _gpu_result = match agent.train_step(explicit_batch) { - Ok(r) => r, - Err(e) => { - let msg = e.to_string(); - if msg.contains("Early stopping") || msg.contains("Gradient collapse") { - return Err(anyhow::anyhow!("{}", msg)); - } - return Err(anyhow::anyhow!("GPU train step FAILED (no CPU fallback): {e}")); - } - }; - - // CUDA: guard kernel accumulates + NaN check - #[cfg(feature = "cuda")] - if let Some(ref mut guard) = self.training_guard { - let gr = guard.check_and_accumulate( - &_gpu_result.loss_gpu, - &_gpu_result.grad_norm_gpu, - 1e6_f32, - guard_collapse_thresh2, - !guard_past_warmup2, - ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?; - if gr.halt_nan { - return Err(anyhow::anyhow!( - "NaN/Inf at step {}: loss={}, grad={}", - train_step_count, gr.raw_loss, gr.raw_grad_norm - )); - } - if gr.halt_grad_collapse { - agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - } - train_step_count += 1; - self.gradient_logging_step += 1; - } - } - } - - // ═══ Single epoch-boundary readback ═══ - if train_step_count > 0 { - let n = train_step_count as f64; - - // CUDA: read from guard's device-resident accumulator - #[cfg(feature = "cuda")] - let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard { - let (al, ag) = guard.read_accumulators() - .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?; - (al as f32, ag as f32) - } else { - (0.0_f32, 0.0_f32) - }; - - #[cfg(not(feature = "cuda"))] - let (avg_loss, avg_grad) = (0.0_f32, 0.0_f32); - - // Q-value estimation at epoch boundary - let avg_q = { - let mut agent = self.agent.write().await; - self.estimate_avg_q_value_with_early_stopping(&mut agent).await? - }; - - // Epoch-level safety checks (same as GPU-PER path) - if !avg_loss.is_finite() || !avg_grad.is_finite() { - let msg = format!( - "SAFETY: NaN/Inf in epoch {} avg — loss={:.6}, grad={:.6}", - epoch, avg_loss, avg_grad - ); - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{} (stopping training)", msg)); - }, - crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { - debug!("{} (continuing training)", msg); - }, - } - } - - // Loss history (epoch-level average) - self.safety_loss_history.push_back(avg_loss); - if self.safety_loss_history.len() > 30 { - self.safety_loss_history.pop_front(); - } - - epoch_loss = avg_loss as f64 * n; - epoch_q_value = avg_q * n; - epoch_gradient_norm = avg_grad as f64 * n; - - // Metrics aggregator - let current_avg_reward = if !monitor.reward_history.is_empty() { - monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 - } else { - 0.0 - }; - self.metrics_aggregator.record(avg_loss, avg_q as f32, current_avg_reward); - self.metrics_aggregator.record_gradient(avg_grad); - - if self.metrics_aggregator.should_log(&self.logging_config) { - let aggregated = self.metrics_aggregator.aggregate_and_clear(); - log_training_progress(&aggregated); - } - - // Gradient diagnostics at epoch boundary - { - let mut agent = self.agent.write().await; - agent.log_diagnostics(avg_grad) - .map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - - monitor.track_q_value_range(avg_q); - } - } // end non-GPU-PER path - - // Phase 2b: Sync GPU weight copies after training updates - #[cfg(feature = "cuda")] - if let Some(ref mut collector) = self.gpu_experience_collector { - let agent = self.agent.read().await; - let dqn_ref: Option<&crate::dqn::DQN> = match &*agent { - DQNAgentType::Standard(ref dqn) => Some(dqn), - DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()), - }; - if let Some(dqn) = dqn_ref { - // Sync online weights: branching > plain dueling > hybrid (distributional+dueling) - // When branching, sync_online_weights internally uses branch_0 key names - // for the advantage slot, so we must pass the branching VarMap. - let online_synced = if let Some(ref bn) = dqn.branching_q_network { - collector.sync_online_weights(bn.vars()).is_ok() - } else if let Some(ref online) = dqn.dueling_q_network { - collector.sync_online_weights(online.vars()).is_ok() - } else if let Some(ref online) = dqn.dist_dueling_q_network { - collector.sync_online_weights(online.vars()).is_ok() - } else { - false - }; - if !online_synced { - return Err(anyhow::anyhow!("GPU online weight sync FAILED -- stale Q-values would corrupt training")); - } - // Sync target weights: branching > plain dueling > hybrid - let target_synced = if let Some(ref bn) = dqn.branching_target_network { - collector.sync_target_weights(bn.vars()).is_ok() - } else if let Some(ref target) = dqn.dueling_target_network { - collector.sync_target_weights(target.vars()).is_ok() - } else if let Some(ref target) = dqn.dist_dueling_target_network { - collector.sync_target_weights(target.vars()).is_ok() - } else { - false - }; - if !target_synced { - return Err(anyhow::anyhow!("GPU target weight sync FAILED -- stale target Q-values would corrupt training")); - } - - // Sync branching DQN extra heads (order + urgency, branches 1+2) - if let Some(ref bn) = dqn.branching_q_network { - collector.sync_online_branching(bn.vars()) - .map_err(|e| anyhow::anyhow!("GPU online branching weight sync FAILED: {e}"))?; - } - if let Some(ref bn) = dqn.branching_target_network { - collector.sync_target_branching(bn.vars()) - .map_err(|e| anyhow::anyhow!("GPU target branching weight sync FAILED: {e}"))?; - } - - // Sync RMSNorm weights for distributional dueling networks (D6) - if let Some(ref online) = dqn.dist_dueling_q_network { - collector.sync_online_rmsnorm(online.vars()) - .map_err(|e| anyhow::anyhow!("GPU online RMSNorm weight sync FAILED: {e}"))?; - } - if let Some(ref target) = dqn.dist_dueling_target_network { - collector.sync_target_rmsnorm(target.vars()) - .map_err(|e| anyhow::anyhow!("GPU target RMSNorm weight sync FAILED: {e}"))?; - } - } - drop(agent); - - // NOTE: curiosity weights are now trained in-place on GPU by - // GpuCuriosityTrainer during experience collection (zero CPU sync). - // The old sync_curiosity_weights_from() call is no longer needed. - } - - // Flush GPU-accumulated max priority to CPU (single readback per epoch - // instead of ~8 per-batch readbacks). Safe because insert_batch uses the - // PREVIOUS max_priority — slightly stale is fine for proportional PER. - #[cfg(feature = "cuda")] - if train_step_count > 0 { - let agent = self.agent.read().await; - if let Err(e) = agent.flush_max_priority() { - debug!("GPU max_priority flush failed (non-fatal): {}", e); - } - } - - let epoch_duration = epoch_start.elapsed(); - - // Calculate epoch metrics (average over training steps, not samples) - let (avg_loss, avg_q_value, avg_grad_norm) = if train_step_count > 0 { - ( - epoch_loss / train_step_count as f64, - epoch_q_value / train_step_count as f64, - epoch_gradient_norm / train_step_count as f64, - ) - } else { - // Early epochs before replay buffer fills - (0.0, 0.0, 0.0) - }; - - total_loss += avg_loss; - total_q_value += avg_q_value; - total_gradient_norm += avg_grad_norm; - - // Calculate average reward for this epoch - let epoch_avg_reward = if !monitor.reward_history.is_empty() { - monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 - } else { - 0.0 - }; - total_reward += epoch_avg_reward as f64; - - // WAVE 30: Log epoch end with aggregated metrics - use crate::dqn::logging::AggregatedMetrics; - let epoch_metrics = AggregatedMetrics { - mean_loss: avg_loss as f32, - std_loss: 0.0, // Not tracked per epoch - mean_q_value: avg_q_value as f32, - std_q_value: 0.0, // Not tracked per epoch - mean_reward: epoch_avg_reward, - mean_gradient_norm: avg_grad_norm as f32, - batch_count: train_step_count, - }; - log_epoch_end(epoch + 1, &epoch_metrics, epoch_duration.as_secs_f64()); - - // Epoch-end: download GPU monitoring summary (48 bytes, one transfer) - #[cfg(feature = "cuda")] - if let Some(ref mon) = self.gpu_monitoring { - if let Ok(summary) = mon.download_summary() { - if summary.total_experiences > 0 { - info!( - "GPU epoch summary: mean_reward={:.6}, std={:.6}, sharpe={:.3}, actions={:?}", - summary.mean_reward, summary.reward_std, summary.sharpe_estimate, - summary.action_counts - ); - // Feed single mean reward into pnl_history for Sharpe-based early stopping. - // Do NOT push N times — that creates zero variance and NaN Sharpe. - // The MonitoringSummary already has the correct Sharpe from the full GPU distribution. - self.pnl_history.push_back(summary.mean_reward as f64); - if self.pnl_history.len() > 1000 { - self.pnl_history.pop_front(); - } - // Feed monitor with summary stats for downstream metrics aggregation - monitor.track_reward(summary.mean_reward); - for (idx, &count) in summary.action_counts.iter().enumerate() { - if count > 0 { - // Route exposure→factored for consistent 45-action monitoring. - // track_action() updates both action_counts and factored_action_counts. - if let Ok(exp_level) = crate::dqn::action_space::ExposureLevel::from_index(idx) { - let factored = self.route_action(exp_level, self.hyperparams.avg_spread as f32); - monitor.track_action(&factored); - } else { - monitor.track_action_by_exposure(idx); - } - } - } - } - } - } - - // VERBOSE: Log reward statistics every 10 epochs - if (epoch + 1) % 10 == 0 && !monitor.reward_history.is_empty() { - let rewards = &monitor.reward_history; - let reward_mean = rewards.iter().sum::() / rewards.len() as f32; - let reward_variance = rewards.iter() - .map(|r| (r - reward_mean).powi(2)) - .sum::() / rewards.len() as f32; - let reward_std = reward_variance.sqrt(); - let reward_min = rewards.iter().copied().fold(f32::INFINITY, f32::min); - let reward_max = rewards.iter().copied().fold(f32::NEG_INFINITY, f32::max); - let non_zero_count = rewards.iter().filter(|&&r| r.abs() > 1e-9).count(); - let non_zero_pct = (non_zero_count as f32 / rewards.len() as f32) * 100.0; - - info!( - "REWARD_STATS: epoch={}, mean={:.6}, std={:.6}, min={:.6}, max={:.6}, non_zero={}/{} ({:.1}%)", - epoch + 1, - reward_mean, - reward_std, - reward_min, - reward_max, - non_zero_count, - rewards.len(), - non_zero_pct - ); - } - - // WAVE 1.2 SAFETY #7: Training Anomaly Detector (plateau detection) - if self.safety_loss_history.len() >= 10 { - let recent_losses: Vec = self.safety_loss_history - .iter() - .rev() - .take(10) - .copied() - .collect(); - - let mean: f32 = recent_losses.iter().sum::() / 10.0; - let variance: f32 = recent_losses - .iter() - .map(|&x| (x - mean).powi(2)) - .sum::() / 10.0; - let std_dev = variance.sqrt(); - - // Alert if loss is stuck (variance < 1% of mean) - if std_dev < mean * 0.01 && mean > 1e-6 { - self.safety_loss_plateau_counter += 1; - - if self.safety_loss_plateau_counter >= 10 { - debug!( - "SAFETY: Training stuck for {} epochs (loss variance: {:.6}, mean: {:.6})", - self.safety_loss_plateau_counter, std_dev, mean - ); - } - } else { - self.safety_loss_plateau_counter = 0; - } - } - - // WAVE 3 AGENT A3: Accumulate action counts for constraint checking - for (i, count) in monitor.action_counts.iter().enumerate() { - total_action_counts[i] += count; - } - for (i, count) in monitor.factored_action_counts.iter().enumerate() { - total_factored_action_counts[i] += count; - } - - // Run monitoring validation at end of epoch - if let Err(e) = monitor.validate_all() { - return Err(e); // Abort training if critical bug detected - } - - // Get current epsilon for logging - let current_epsilon = self.get_epsilon().await?; - - // Get Q-value range statistics - let (q_min, q_max, q_mean) = monitor.get_q_value_stats(); - - info!( - "Epoch {}/{}: train_loss={:.6}, Q-value={:.4}, grad_norm={:.6}, train_steps={}, epsilon={:.4}, duration={:.2}s", - epoch + 1, - self.hyperparams.epochs, - avg_loss, - avg_q_value, - avg_grad_norm, - train_step_count, - current_epsilon, - epoch_duration.as_secs_f64() - ); - - // BUG #29 FIX: Update epsilon once per epoch (not per batch) - // Skip epsilon decay when noisy nets are enabled — epsilon is fixed at noisy_epsilon_floor - if !self.hyperparams.use_noisy_nets { - let mut agent = self.agent.write().await; - agent.update_epsilon(); - } - - // Anneal noisy sigma schedule at epoch boundary (when enabled) - if let Some(ref mut scheduler) = self.noisy_sigma_scheduler { - // Step the scheduler by the number of training steps in this epoch - for _ in 0..train_step_count { - scheduler.step(); - } - let sigma_scale = scheduler.get_sigma(); - let mut agent = self.agent.write().await; - agent.set_noise_sigma_scale(sigma_scale); - } - - // Reset count-based exploration bonus at epoch boundary - { - let mut agent = self.agent.write().await; - agent.reset_count_bonus(); - } - - // M2: Refresh stale PER priorities at epoch boundary - // With N-step returns, priorities computed at insertion become stale as the - // network trains. Refresh the oldest 5% of the buffer per epoch. - if self.hyperparams.use_per && self.hyperparams.n_steps > 1 { - let agent = self.agent.read().await; - let buffer = agent.memory(); - let buf_len = buffer.len(); - let refresh_limit = (buf_len / 20).max(32).min(256); // 5% of buffer, capped at 256 - let max_age = 500; // Refresh entries not updated for 500+ training steps - let stale_indices = buffer.get_stale_indices(max_age, refresh_limit); - if !stale_indices.is_empty() { - // Get experiences and recompute approximate TD errors via Q-value magnitude - let experiences = buffer.get_experiences_at(&stale_indices); - // Batch all valid states for a single GPU round-trip instead of - // per-experience {from_slice + forward + to_scalar}. - let mut flat_states: Vec = Vec::with_capacity(stale_indices.len() * 64); - let mut valid_indices = Vec::with_capacity(stale_indices.len()); - let mut state_dim: Option = None; - - for (idx, exp_opt) in stale_indices.iter().zip(experiences.iter()) { - if let Some(Some(exp)) = exp_opt.as_ref().map(Some) { - if let Some(sd) = state_dim { - if exp.state.len() != sd { - continue; // skip mismatched dimensions - } - } else { - state_dim = Some(exp.state.len()); - } - flat_states.extend_from_slice(&exp.state); - valid_indices.push(*idx); - } - } - - if let Some(sd) = state_dim { - if !valid_indices.is_empty() { - let batch_size = valid_indices.len(); - let td_result = candle_core::Tensor::from_vec( - flat_states, &[batch_size, sd], agent.device(), - ) - .and_then(|bt| { - agent.forward(&bt).map_err(|e| { - candle_core::Error::Msg(format!("forward: {e}")) - }) - }) - .and_then(|q_vals| q_vals.max(candle_core::D::Minus1)) - .and_then(|mq| { - // GPU-side: abs + clamp(min=0.01) + NaN→0.01 - // NaN detection: NaN != NaN, so ne(self) finds NaN entries - let abs_q = mq.abs()?; - let floor = candle_core::Tensor::new(0.01_f32, abs_q.device())?; - let clamped = abs_q.broadcast_maximum(&floor)?; - let nan_mask = clamped.ne(&clamped)?; // true where NaN - let floor_bcast = floor.broadcast_as(clamped.shape())?; - let safe = nan_mask.where_cond(&floor_bcast, &clamped)?; - Ok(safe) - }); - - if let Ok(td_errors_gpu) = td_result { - // GPU PER: update priorities without CPU readback - #[cfg(feature = "cuda")] - { - let idx_u32: Vec = valid_indices.iter() - .map(|&i| i as u32) - .collect(); - if let Ok(idx_tensor) = candle_core::Tensor::new( - idx_u32, agent.device(), - ) { - if let Err(e) = agent.update_priorities_gpu( - &idx_tensor, &td_errors_gpu, - ) { - debug!("M2: GPU priority refresh failed (non-fatal): {}", e); - } else { - debug!( - "Epoch {}: Refreshed {} stale PER priorities (max_age={}, GPU)", - epoch + 1, valid_indices.len(), max_age - ); - } - } - } - #[cfg(not(feature = "cuda"))] - { - // CPU-only: scalar readback per priority (no to_vec1) - let mut td_vec = Vec::with_capacity(valid_indices.len()); - for i in 0..valid_indices.len() { - if let Ok(v) = td_errors_gpu.narrow(0, i, 1) - .and_then(|t| t.to_scalar::()) { - td_vec.push(v); - } - } - if td_vec.len() == valid_indices.len() { - if let Err(e) = buffer.update_priorities(&valid_indices, &td_vec) { - debug!("M2: Priority refresh failed (non-fatal): {}", e); - } - } - } - } - } - } - } - } - - // WAVE 26 P0.6: Update learning rate with scheduler (warmup + decay) - self.lr_scheduler.step(); - let current_lr = self.lr_scheduler.get_lr(); - if epoch % 10 == 0 { - info!( - "Learning rate scheduled update: epoch={}, lr={:.2e} (initial={:.2e})", - epoch + 1, - current_lr, - self.lr_scheduler.get_initial_lr() - ); - } - // Apply scheduled LR to the optimizer - let initial_lr = self.lr_scheduler.get_initial_lr(); - if initial_lr > 0.0 { - let decay_factor = current_lr / initial_lr; - self.agent.write().await.update_learning_rate(decay_factor)?; - } - training_metrics::set_learning_rate("dqn", "current", current_lr); - - // WAVE 9-11: Log Q-value range for production monitoring - if train_step_count > 0 { - info!( - "Epoch {}/{}: Q-value range=[{:.2}, {:.2}], mean={:.2}", - epoch + 1, - self.hyperparams.epochs, - q_min, - q_max, - q_mean - ); - - // WAVE 9-11: Warning threshold (500K as per production test report) - const Q_VALUE_WARNING_THRESHOLD: f64 = 500_000.0; - if q_max > Q_VALUE_WARNING_THRESHOLD { - warn!( - "⚠️ Q-value explosion detected at epoch {}: max Q-value {:.2e} exceeds threshold {:.2e}", - epoch + 1, - q_max, - Q_VALUE_WARNING_THRESHOLD - ); - warn!("Consider:"); - warn!(" • Reducing learning rate (current: {:.2e})", self.hyperparams.learning_rate); - warn!(" • Enabling target network soft updates (Polyak averaging, tau=0.005)"); - warn!(" • Adjusting reward scaling"); - training_metrics::record_gradient_explosion("dqn", "current"); - } - - // M3: Combined Q-value diagnostics (gap + per-action averages) - // Single forward pass + readback instead of two separate ones. - if let Some(((gap_mean, gap_min, gap_max), per_action_avgs)) = - self.compute_epoch_q_diagnostics().await - { - info!( - "Epoch {}/{}: Q-value gap (best-2nd): mean={:.4}, min={:.4}, max={:.4}", - epoch + 1, self.hyperparams.epochs, - gap_mean, gap_min, gap_max - ); - - let names = ["S100", "S50", "Flat", "L50", "L100"]; - let parts: Vec = names - .iter() - .zip(per_action_avgs.iter()) - .map(|(n, q)| format!("{}={:.4}", n, q)) - .collect(); - info!( - "Epoch {}/{}: Per-action Q: {}", - epoch + 1, self.hyperparams.epochs, - parts.join(", ") - ); - } - } - - // Compute validation loss - let val_loss = self.compute_validation_loss().await?; - info!( - "Epoch {}/{}: val_loss={} (backtest Sharpe proxy)", - epoch + 1, - self.hyperparams.epochs, - if val_loss.abs() < 1e-10 { "N/A".to_owned() } else { format!("{val_loss:.6}") } - ); - - // Per-epoch Prometheus metrics for monitoring service - training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64); - training_metrics::set_epoch_loss("dqn", "current", avg_loss); - training_metrics::set_validation_loss("dqn", "current", val_loss); - if epoch_duration.as_secs_f64() > 0.0 { - training_metrics::set_batches_per_second( - "dqn", "current", - train_step_count as f64 / epoch_duration.as_secs_f64(), - ); - } - // Tier 1: RL diagnostics + training health - training_metrics::set_q_value_stats("dqn", "current", q_mean, q_max); - training_metrics::set_gradient_norm("dqn", "current", avg_grad_norm); - training_metrics::set_epoch_duration("dqn", "current", epoch_duration.as_secs_f64()); - { - let agent = self.agent.read().await; - if let Ok(buf_size) = agent.get_replay_buffer_size() { - training_metrics::set_replay_buffer_size("dqn", "current", buf_size as f64); - } - } - - // Adaptive tau: detect Q-value overestimation and increase Polyak averaging - if epoch > 0 { - let q_mean_growth = q_mean - self.prev_epoch_q_mean; - // After percentage-reward fix (denormalized portfolio), Q-values are - // on a ~0.01/step scale, not raw dollars. Old thresholds (0.5 / 0.1) - // would never trigger. - if q_mean_growth > 0.005 { - self.adaptive_tau = (self.adaptive_tau * 2.0).min(0.01); - warn!( - "Q-value drift detected (Δ={:.3}), increasing tau to {:.4}", - q_mean_growth, self.adaptive_tau - ); - } else if q_mean_growth < 0.001 { - self.adaptive_tau = (self.adaptive_tau * 0.9).max(self.hyperparams.tau); - } else { - // Q-value growth in normal range — keep current tau - } - // Propagate adaptive tau to agent config - { - let mut agent = self.agent.write().await; - match &mut *agent { - DQNAgentType::Standard(dqn) => dqn.config.tau = self.adaptive_tau, - DQNAgentType::RegimeConditional(regime) => { - regime.primary_head_mut().config.tau = self.adaptive_tau; - } - } - } - } - self.prev_epoch_q_mean = q_mean; - - // WAVE 9-11 PRODUCTION: Track action diversity per epoch - // Calculate active factored actions (45-action space when branching) - let epoch_total_factored: usize = monitor.factored_action_counts.iter().sum(); - let epoch_total_exposure: usize = monitor.action_counts.iter().sum(); - let (active_actions_count, action_space_size, diversity_percentage) = - if epoch_total_factored > 0 { - // Branching DQN: report factored 45-action diversity - let active_threshold = (epoch_total_factored as f64 * 0.005).max(1.0); - let active = monitor.factored_action_counts.iter() - .filter(|&&count| count as f64 >= active_threshold) - .count(); - (active, 45_usize, (active as f64 / 45.0) * 100.0) - } else { - // Non-branching: report 5-action exposure diversity - let active_threshold = (epoch_total_exposure as f64 * 0.005).max(1.0); - let active = monitor.action_counts.iter() - .filter(|&&count| count as f64 >= active_threshold) - .count(); - (active, 5_usize, (active as f64 / 5.0) * 100.0) - }; - - // Log action diversity - info!( - "Epoch {}/{}: Action diversity={}/{} ({:.1}%)", - epoch + 1, - self.hyperparams.epochs, - active_actions_count, - action_space_size, - diversity_percentage - ); - - // Compute normalized entropy of epoch action distribution - let (epoch_entropy, entropy_total, entropy_size) = if epoch_total_factored > 0 { - // Use factored 45-action distribution for entropy - let mut entropy_raw = 0.0_f64; - for &count in monitor.factored_action_counts.iter() { - if count > 0 { - let p = count as f64 / epoch_total_factored as f64; - entropy_raw -= p * p.ln(); - } - } - (entropy_raw / 45.0_f64.ln(), epoch_total_factored, 45_usize) - } else if epoch_total_exposure > 0 { - let mut entropy_raw = 0.0_f64; - for &count in monitor.action_counts.iter() { - if count > 0 { - let p = count as f64 / epoch_total_exposure as f64; - entropy_raw -= p * p.ln(); - } - } - (entropy_raw / 5.0_f64.ln(), epoch_total_exposure, 5_usize) - } else { - (0.0, 0, 5) - }; - let _ = (entropy_total, entropy_size); // used implicitly by entropy value - - info!( - " Exploration: entropy={:.3} (1.0=uniform), epsilon={:.4}, noisy_nets={}, count_bonus={}", - epoch_entropy, - self.get_epsilon().await.unwrap_or(0.0), - self.hyperparams.use_noisy_nets, - self.hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, - ); - - // Emit exploration diagnostics to Prometheus - training_metrics::set_action_entropy("dqn", "current", epoch_entropy); - training_metrics::set_action_diversity("dqn", "current", diversity_percentage / 100.0); - - // Warning if diversity drops below 40% (2 of 5 exposure levels) - const DIVERSITY_THRESHOLD: usize = 2; // 40% of 5 exposure levels - if active_actions_count < DIVERSITY_THRESHOLD { - warn!( - "⚠️ LOW ACTION DIVERSITY: {}/5 exposure levels (<40%), entropy={:.3}", - active_actions_count, epoch_entropy, - ); - } - - // WAVE P2: Log episode statistics - let (mean_len, std_len, min_len, max_len, exit_counts) = monitor.get_episode_stats(); - if !monitor.episode_lengths.is_empty() { - let total_episodes = monitor.episode_lengths.len(); - - info!( - "WAVE P2 Episode Stats [Epoch {}]: {} episodes, length: mean={:.1}±{:.1}, min={}, max={}", - epoch + 1, - total_episodes, - mean_len, - std_len, - min_len, - max_len - ); - - info!( - " Exit breakdown: profit={}({:.1}%), stop={}({:.1}%), time={}({:.1}%), boundary={}({:.1}%)", - exit_counts[0], (exit_counts[0] as f64 / total_episodes as f64) * 100.0, - exit_counts[1], (exit_counts[1] as f64 / total_episodes as f64) * 100.0, - exit_counts[2], (exit_counts[2] as f64 / total_episodes as f64) * 100.0, - exit_counts[3], (exit_counts[3] as f64 / total_episodes as f64) * 100.0, - ); - } - - // WAVE 3.11: Calculate and log VaR/CVaR from PnL history - if self.pnl_history.len() > 20 { - // Normalize raw dollar PnL to percentage returns relative to initial capital - let capital = self.hyperparams.initial_capital as f64; - let returns: Vec = if capital > 0.0 { - self.pnl_history.iter().map(|&pnl| pnl / capital).collect() - } else { - self.pnl_history.iter().copied().collect() - }; - - // Calculate VaR/CVaR at 95% and 99% confidence levels - // confidence_level=0.05 means we're looking at the worst 5% of returns (95% VaR) - let (var_95, cvar_95) = calculate_var_cvar(&returns, 0.05); - let (var_99, cvar_99) = calculate_var_cvar(&returns, 0.01); - - info!( - "Epoch {}/{}: Risk Metrics - VaR(95%)={:.4}%, CVaR(95%)={:.4}%, VaR(99%)={:.4}%, CVaR(99%)={:.4}% (from {} PnL samples)", - epoch + 1, - self.hyperparams.epochs, - var_95 * 100.0, - cvar_95 * 100.0, - var_99 * 100.0, - cvar_99 * 100.0, - returns.len() - ); - } - - // Epoch financial metrics for monitoring service - // C4: Extract epoch_sharpe for Sharpe-based early stopping & best-checkpoint - let epoch_sharpe = { - let financials = compute_epoch_financials( - &self.pnl_history, - &monitor.action_counts, - 100_000.0, - ); - training_metrics::set_epoch_financial_metrics( - "dqn", "current", - financials.sharpe, - financials.sortino, - financials.win_rate, - financials.max_drawdown, - financials.profit_factor, - financials.total_return, - financials.avg_return, - financials.total_trades as f64, - ); - training_metrics::set_epoch_action_distribution( - "dqn", "current", - financials.buy_pct, - financials.sell_pct, - financials.hold_pct, - ); - info!( - "Epoch {}/{}: Sharpe={:.2} WinRate={:.1}% MaxDD={:.3}% PF={:.2} Return={:+.2}% Trades={}", - epoch + 1, self.hyperparams.epochs, - financials.sharpe, financials.win_rate * 100.0, - financials.max_drawdown * 100.0, financials.profit_factor, - financials.total_return * 100.0, financials.total_trades, - ); - // QuestDB: persist epoch row for historical queries in Grafana - { - let bps = if epoch_duration.as_secs_f64() > 0.0 { - train_step_count as f64 / epoch_duration.as_secs_f64() - } else { - 0.0 - }; - let buf_size = { - let agent = self.agent.read().await; - agent.get_replay_buffer_size().unwrap_or(0) as f64 - }; - let run_id = std::env::var("HOSTNAME").unwrap_or_else(|_| "local".to_owned()); - questdb_sink::record_training_epoch(&questdb_sink::EpochRecord { - model: "dqn", - fold: "current", - run_id: &run_id, - symbol: "", - epoch: (epoch + 1) as u32, - loss: avg_loss, - val_loss, - batches_per_second: bps, - epoch_duration_secs: epoch_duration.as_secs_f64(), - q_mean, - q_max, - gradient_norm: avg_grad_norm, - learning_rate: current_lr, - replay_buffer_size: buf_size, - action_entropy: epoch_entropy, - action_diversity: diversity_percentage / 100.0, - sharpe: financials.sharpe, - sortino: financials.sortino, - win_rate: financials.win_rate, - max_drawdown: financials.max_drawdown, - profit_factor: financials.profit_factor, - total_return: financials.total_return, - avg_return: financials.avg_return, - total_trades: financials.total_trades as f64, - }); - } - - financials.sharpe - }; - - // Track metrics for early stopping - self.loss_history.push(avg_loss); - self.q_value_history.push(avg_q_value); - self.val_loss_history.push(val_loss); - self.sharpe_history.push(epoch_sharpe); - - // MEMORY LEAK FIX: Limit history vectors to prevent unbounded growth - // Keep last 100 epochs (sufficient for early stopping window of 5) - const MAX_HISTORY_LEN: usize = 100; - if self.loss_history.len() > MAX_HISTORY_LEN { - self.loss_history.drain(0..50); // Remove oldest 50, keep newest 50 - } - if self.q_value_history.len() > MAX_HISTORY_LEN { - self.q_value_history.drain(0..50); - } - if self.val_loss_history.len() > MAX_HISTORY_LEN { - self.val_loss_history.drain(0..50); - } - if self.sharpe_history.len() > MAX_HISTORY_LEN { - self.sharpe_history.drain(0..50); - } - - // C4 FIX: Save best model checkpoint when Sharpe improves (not val-loss). - // Sharpe directly measures trading quality — the metric we actually optimize. - if train_step_count > 0 && epoch_sharpe > self.best_sharpe { - self.best_sharpe = epoch_sharpe; - self.best_val_loss = val_loss; // Track for logging only - self.best_epoch = epoch + 1; - - info!( - "🎉 New best Sharpe: {:.4} at epoch {} (val_loss={:.6})", - epoch_sharpe, - epoch + 1, - val_loss, - ); - - // WAVE 1.2 SAFETY #8: Checkpoint Verification (before save) - let checkpoint_data = self.serialize_model().await?; - - // Verify checkpoint integrity (check for NaN/Inf in serialized data) - if self.safety_level != crate::safety::SafetyLevel::Permissive { - // Simple check: ensure checkpoint data is not empty and doesn't contain obvious corruption markers - if checkpoint_data.is_empty() { - let msg = "SAFETY: Checkpoint verification failed - empty checkpoint data"; - match self.safety_level { - crate::safety::SafetyLevel::Strict => { - return Err(anyhow::anyhow!("{}", msg)); - }, - crate::safety::SafetyLevel::Normal => { - debug!("{} (continuing anyway)", msg); - }, - _ => {}, - } - } else { - debug!("SAFETY: Checkpoint verification passed ({} bytes)", checkpoint_data.len()); - } - } - - let ckpt_size = checkpoint_data.len() as f64; - let ckpt_start = std::time::Instant::now(); - let best_checkpoint_path = checkpoint_callback( - epoch + 1, - checkpoint_data, - true, // is_best flag - ) - .context("Failed to save best checkpoint")?; - training_metrics::record_checkpoint_save("dqn", "current", ckpt_start.elapsed().as_secs_f64(), ckpt_size); - - info!("Best model saved to: {}", best_checkpoint_path); - } - - // Early stopping checks (skip if no training occurred) - if train_step_count > 0 { - // C4 FIX: Sharpe-based early stopping instead of val-loss. - // check_early_stopping uses Sharpe plateau; patience uses -Sharpe (lower=better API). - let old_should_stop = self.check_early_stopping(avg_q_value, epoch); - let patience_should_stop = if self.hyperparams.early_stopping_enabled - && epoch + 1 >= self.hyperparams.min_epochs_before_stopping { - // Negate Sharpe: EarlyStopping expects "lower is better" - self.early_stopping.should_stop(-epoch_sharpe) - } else { - false - }; - - // Stop if EITHER old criteria OR new patience-based criteria trigger - if let Some(stop_reason) = old_should_stop { - warn!( - "Early stopping triggered at epoch {}/{}: {}", - epoch + 1, - self.hyperparams.epochs, - stop_reason - ); - info!( - "Final metrics: loss={:.6}, Q-value={:.4}", - avg_loss, avg_q_value - ); - - // WAVE 13-A2: Save checkpoint for early stopping (use is_best=false for proper naming) - let checkpoint_data = self - .serialize_model() - .await - .context("Failed to serialize model for early stopping checkpoint")?; - let checkpoint_size = checkpoint_data.len(); - let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false) - .context("Failed to save early stopping checkpoint")?; - info!( - "Early stopping checkpoint saved to: {} ({} bytes)", - checkpoint_path, checkpoint_size - ); - - // WAVE 23 P0 FIX: Return error instead of Ok(metrics) to terminate with non-zero exit code - // This ensures hyperopt can properly detect and kill failing trials early - // The checkpoint has already been saved above, so the model state is preserved - return Err(anyhow::anyhow!( - "Training terminated by early stopping at epoch {}/{}: {}", - epoch + 1, - self.hyperparams.epochs, - stop_reason - )); - } - - // WAVE 24 (Agent 17): New patience-based early stopping check - if patience_should_stop { - warn!( - "🛑 WAVE 24 Early stopping (patience) triggered at epoch {}/{}! No improvement for {} epochs", - epoch + 1, self.hyperparams.epochs, self.hyperparams.gradient_collapse_patience - ); - info!("Best Sharpe: {:.4} at epoch {} (best val_loss proxy: {:.6})", - self.best_sharpe, self.best_epoch, self.early_stopping.get_best_val_loss()); - - // Save checkpoint before stopping - let checkpoint_data = self.serialize_model().await?; - let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false)?; - info!("Patience-based early stopping checkpoint saved to: {}", checkpoint_path); - - // Return error to terminate training (consistent with existing early stopping behavior) - return Err(anyhow::anyhow!("Training terminated by patience-based early stopping at epoch {}", epoch + 1)); - } - } - - // WAVE 13-A2: Save periodic checkpoint every N epochs - if (epoch + 1) % self.hyperparams.checkpoint_frequency == 0 { - info!( - "💾 Saving periodic checkpoint at epoch {}/{}", - epoch + 1, - self.hyperparams.epochs - ); - - let checkpoint_data = self.serialize_model().await?; - let checkpoint_size = checkpoint_data.len(); - let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false) - .context("Failed to save periodic checkpoint")?; - - info!( - "✅ Periodic checkpoint saved: {} ({} bytes)", - checkpoint_path, checkpoint_size - ); - } - } - - let training_duration = start_time.elapsed(); - - // Calculate final metrics - let metrics = self - .create_final_metrics( - total_loss, - total_q_value, - total_gradient_norm, - total_reward, - self.hyperparams.epochs, - training_duration, - false, - total_action_counts, // WAVE 3 AGENT A3 - total_factored_action_counts, - ) - .await?; - - // Update stored metrics - { - let mut stored_metrics = self.metrics.write().await; - *stored_metrics = metrics.clone(); - } - - info!( - "Training completed in {:.2}s: final_loss={:.6}, avg_q_value={:.4}", - training_duration.as_secs_f64(), - metrics.loss, - metrics - .additional_metrics - .get("avg_q_value") - .unwrap_or(&0.0) - ); - - info!("Best model summary:"); - info!( - " Best Sharpe: {:.4} at epoch {} (val_loss={:.6})", - self.best_sharpe, self.best_epoch, self.best_val_loss - ); - info!(" Best model checkpoint: best_model.safetensors"); - - Ok(metrics) - } - - - /// Convert feature vector to TradingState (42 market features → 45-dim state with portfolio) - /// - /// CRITICAL BUG FIX: Features 0-3 are LOG RETURNS (signed), not raw prices. - /// Using .abs() destroys directional information (bullish vs bearish). - /// We now use TradingState::from_normalized() to preserve sign information. - /// - /// Feature mapping: - /// - Features 0-3: OHLC log returns → price_features (signed, normalized) - /// - Features 4-224: All other features → technical_indicators (221 features including Wave D) - /// - /// # Arguments - /// - /// * `feature_vec` - 42-dimensional FeatureVector from extraction pipeline - /// * `close_price` - Current close price for portfolio feature calculation (optional) - /// - /// # Bug #4 Fix - /// - /// Added close_price parameter to enable portfolio feature population from PortfolioTracker. - fn feature_vector_to_state( - &self, - feature_vec: &FeatureVector, - close_price: Option, - ) -> Result { - self.feature_vector_to_state_with_ofi(feature_vec, close_price, None) - } - - fn feature_vector_to_state_with_ofi( - &self, - feature_vec: &FeatureVector, - close_price: Option, - ofi_index: Option, - ) -> Result { - // States are pre-normalized during data loading - let normalized_features: Vec = feature_vec.iter().map(|&v| v as f32).collect(); - - // Features 0-3 are LOG RETURNS - preserve sign information for price direction - let price_features: Vec = vec![ - normalized_features[0], // open log return (can be negative) - normalized_features[1], // high log return (can be negative) - normalized_features[2], // low log return (can be negative) - normalized_features[3], // close log return (can be negative) - ]; - - // 42-FEATURE ARCHITECTURE: Extract market features (indices 4-41) - assert_eq!( - normalized_features.len(), - 42, - "Expected 42 market features (got {})", - normalized_features.len() - ); - let market_features: Vec = normalized_features[4..42] - .iter() - .map(|&x| x as f32) - .collect(); - - // Legacy technical_indicators (empty for 42-feature architecture) - let technical_indicators = vec![]; - - // BUG #36 FIX: Use NORMALIZED portfolio features to prevent Q-value explosion - let portfolio_features = if let Some(price) = close_price { - let price_f32 = price.to_f32().unwrap_or(0.0); - self.portfolio_tracker - .get_portfolio_features(price_f32) - .to_vec() - } else { - vec![0.0, 0.0, 0.0] // Fallback if no price provided - }; - - // OFI regime features: 8 features from MBP-10 order book data. - // When OFI is enabled (mbp10_data_dir set), always return 8 features - // (zeros if data didn't load) to match state_dim=53. - let ofi_enabled = self.hyperparams.mbp10_data_dir.is_some(); - let regime_features: Vec = if let (Some(ofi), Some(idx)) = (&self.ofi_features, ofi_index) { - ofi.get(idx) - .map(|f| f.iter().map(|&v| v as f32).collect()) - .unwrap_or_else(|| vec![0.0; 8]) - } else if ofi_enabled { - vec![0.0; 8] - } else { - vec![] - }; - - // Use from_normalized() to preserve sign information - Ok(TradingState::from_normalized( - price_features, - technical_indicators, - market_features, - portfolio_features, - regime_features, - )) - } - - /// Select action using epsilon-greedy - async fn select_action(&self, state: &TradingState) -> Result { - let _agent = self.agent.read().await; - - // Convert state to tensor with tensor core alignment padding - let state_vec = state.to_vector(); - let raw_dim = state_vec.len(); - let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - let padded: Vec = if aligned_dim > raw_dim { - let mut v = state_vec.to_vec(); - v.resize(aligned_dim, 0.0); - v - } else { - state_vec.to_vec() - }; - let state_tensor = Tensor::new(&*padded, &self.device) - .map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))? - .unsqueeze(0)?; // Add batch dimension - - // Get Q-values (epsilon-greedy handled by agent internally) - let action_idx = self.epsilon_greedy_action(&state_tensor).await?; - - let exposure = ExposureLevel::from_index(action_idx) - .map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?; - // Phase C: Use smart routing with trainer's spread/vol EMAs - Ok(self.route_action(exposure, self.hyperparams.avg_spread as f32)) - } - - /// Phase C: Route an exposure-level action using smart order routing. - /// - /// Uses the trainer's running spread/volatility EMAs to determine - /// optimal order type (Market/Limit/IoC) and urgency (Patient/Normal/Aggressive). - /// The DQN selects exposure; OrderRouter selects execution strategy. - fn route_action(&self, exposure: ExposureLevel, current_spread: f32) -> FactoredAction { - OrderRouter::route( - exposure, - current_spread, - self.hyperparams.avg_spread as f32, - self.vol_ema as f32, - self.median_vol as f32, - ) - } - - /// Phase C: Simulate order fill and return result. - /// - /// Returns (action, fill_result) — if not filled, action is overridden to Flat - /// so the agent learns that limit orders in certain conditions don't execute. - fn simulate_fill( - &self, - action: FactoredAction, - step: usize, - ) -> (FactoredAction, FillResult) { - let normalized_vol = if self.median_vol > 0.0 { - (self.vol_ema / self.median_vol) as f32 - } else { - 1.0 - }; - let spread_bps = self.hyperparams.avg_spread * 10000.0; // fractional → bps - - let fill_result = self.fill_simulator.simulate_fill( - action.order, - action.urgency, - normalized_vol, - spread_bps, - step, - action.exposure as usize, - ); - - if fill_result.filled { - (action, fill_result) - } else { - // Order didn't fill — position stays unchanged (Flat action, no trade) - (OrderRouter::route_default(ExposureLevel::Flat), fill_result) - } - } - - /// Select actions for a batch of states (GPU-optimized) - /// - /// This method reduces GPU kernel launches by batching all action selections - /// into a single forward pass. Provides 125× reduction in kernel launches - /// compared to sequential select_action() calls. - /// - /// # Performance Impact - /// - Single GPU kernel launch for entire batch (vs. one per sample) - /// - Reduced CPU-GPU synchronization overhead - /// - Better GPU utilization through larger batch sizes - /// - /// # Arguments - /// * `states` - Slice of TradingState objects to process - /// - /// # Returns - /// Vector of TradingAction decisions (same order as input states) - async fn select_actions_batch(&mut self, states: &[TradingState]) -> Result> { - if states.is_empty() { - return Ok(Vec::new()); - } - - let agent = self.agent.read().await; - let batch_size = states.len(); - - // Get state dimension from first state, then align for tensor cores - let first_vec = states - .first() - .map(|s| s.to_vector()) - .ok_or_else(|| anyhow::anyhow!("Empty states slice"))?; - let raw_state_dim = first_vec.len(); - let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &self.device); - let pad = aligned_dim - raw_state_dim; - - // Pre-allocate flat buffer, zero-padding each state to aligned dimension - let mut flat_states = Vec::with_capacity(batch_size * aligned_dim); - flat_states.extend_from_slice(&first_vec); - flat_states.extend(std::iter::repeat_n(0.0_f32, pad)); - - for (i, state) in states.iter().enumerate().skip(1) { - let vec = state.to_vector(); - if vec.len() != raw_state_dim { - return Err(anyhow::anyhow!( - "State {} dimension mismatch: expected {}, got {}", - i, - raw_state_dim, - vec.len() - )); - } - flat_states.extend_from_slice(&vec); - flat_states.extend(std::iter::repeat_n(0.0_f32, pad)); - } - - // Create batched tensor directly from flat buffer - let batch_tensor = Tensor::from_vec(flat_states, (batch_size, aligned_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))?; - - // FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor. - // Previously used get_epsilon() which returns the decayed epsilon (0.0 with noisy nets), - // making the batch path have ZERO random exploration — root cause of action collapse. - let base_epsilon = agent.get_effective_epsilon() as f64; - let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon); - let epsilon = adjusted_epsilon as f32; - - debug!("Epsilon: base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon); - - // Single forward pass for all samples (GPU-optimized). - // For RegimeConditional, forward() now blends all 3 heads via regime masks. - let batch_q_values = agent - .forward(&batch_tensor) - .map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?; - - // Branching DQN: get per-branch Q-values while agent lock is held. - // Returns (exposure [batch,5], order [batch,3], urgency [batch,3]). - let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = - if self.hyperparams.use_branching { - agent - .batch_branching_q_values(&batch_tensor) - .map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))? - } else { - None - }; - - drop(agent); // Release lock early - - #[cfg(not(feature = "cuda"))] - { - let _ = (&batch_q_values, &branching_q_tensors, epsilon, batch_size); - #[allow(clippy::needless_return)] - return Err(anyhow::anyhow!("Batch action selection requires CUDA — enable the `cuda` feature")); - } - - // Fused GPU epsilon-greedy: argmax + RNG in a single CUDA kernel launch. - // Eliminates the intermediate GPU→CPU sync from Candle's argmax(). - // Lazy-init the GPU action selector on first call - #[cfg(feature = "cuda")] - { - if self.gpu_action_selector.is_none() && self.device.is_cuda() { - let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( - &self.device, - self.hyperparams.batch_size.max(batch_size).max(8192), - 0xDEAD_BEEF_CAFE_u64, - ).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?; - info!("GPU action selector initialized for select_actions_batch"); - self.gpu_action_selector = Some(selector); - } - - let selector = self.gpu_action_selector.as_mut() - .ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?; - - // Branching path: per-branch epsilon-greedy → factored indices (0-44) - let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { - selector - .select_actions_branching(q_exp, q_ord, q_urg, epsilon) - .map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))? - } else { - // Non-branching: exposure-only epsilon-greedy → GPU route to factored - let exposure_tensor = selector - .select_actions(&batch_q_values, epsilon, batch_size, 5) - .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; - - selector.route_exposure_to_factored( - &exposure_tensor, - batch_size, - self.hyperparams.avg_spread as f32, - self.hyperparams.avg_spread as f32, - self.vol_ema as f32, - self.median_vol as f32, - ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? - }; - - // Per-element scalar readback (no to_vec1) - // narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar - let mut actions = Vec::with_capacity(batch_size); - for i in 0..batch_size { - let idx = factored_tensor.narrow(0, i, 1) - .and_then(|t| t.squeeze(0)) - .and_then(|t| t.to_scalar::()) - .map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?; - let action = FactoredAction::from_index(idx as usize) - .map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?; - actions.push(action); - } - Ok(actions) - } - } - - /// GPU-optimized batch action selection using pre-built state tensor. - /// - /// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch). - /// GPU-batched action selection with optional fused routing + fill simulation. - /// - /// Returns `(actions, gpu_handled_fill)`: - /// - `gpu_handled_fill = true`: actions are post-fill (routed + fill-checked by GPU kernel). - /// Caller must NOT apply CPU `route_action()` or `simulate_fill()` — already done. - /// - `gpu_handled_fill = false`: actions have basic routing only. Caller should apply - /// CPU routing + fill as before. - /// - /// The fused kernel (`epsilon_greedy_routed`) is used when: - /// 1. GPU action selector is available (CUDA device) - /// 2. Not using branching DQN (branching learns order type via network heads) - /// 3. Median volatility > 0 (fill simulation requires vol context) - #[cfg(feature = "cuda")] - async fn select_actions_batch_gpu( - &mut self, - batch_tensor: &Tensor, - batch_start: usize, - ) -> Result<(Vec, bool)> { - let batch_size = batch_tensor.dims()[0]; - if batch_size == 0 { - return Ok((Vec::new(), false)); - } - - let agent = self.agent.read().await; - - // FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor. - // Same fix as select_actions_batch() — previously used get_epsilon() which - // returned 0.0 with noisy nets, causing zero exploration in GPU batch path. - let base_epsilon = agent.get_effective_epsilon() as f64; - let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon); - let epsilon = adjusted_epsilon as f32; - debug!("Epsilon (GPU path): base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon); - - // Single forward pass — tensor already on GPU, no construction needed - let batch_q_values = agent - .forward(batch_tensor) - .map_err(|e| anyhow::anyhow!("GPU batched forward pass failed: {}", e))?; - - // C2 FIX: Count bonus removed from Q-value computation. - // Noisy nets are the sole exploration mechanism during training. - // Count bonus kept for diversity metrics only (record_action tracking). - - // Branching DQN: get per-branch Q-values if branching mode is active. - // Uses unified dispatch that supports both Standard and RegimeConditional. - // For RC, regime classification masks blend per-branch Q-values from all 3 heads. - #[cfg(feature = "cuda")] - let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = if self.hyperparams.use_branching { - agent - .batch_branching_q_values(batch_tensor) - .map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))? - } else { - None - }; - - drop(agent); - - // Fused GPU epsilon-greedy: argmax + RNG + routing in CUDA kernels. - // Lazy-init GPU action selector on first call. - if self.gpu_action_selector.is_none() && self.device.is_cuda() { - let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( - &self.device, - self.hyperparams.batch_size.max(batch_size).max(8192), - 0xDEAD_BEEF_CAFE_u64, - ).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?; - info!("GPU action selector initialized for select_actions_batch_gpu"); - self.gpu_action_selector = Some(selector); - } - - let selector = self.gpu_action_selector.as_mut() - .ok_or_else(|| anyhow::anyhow!("GPU action selector not initialized on non-CUDA device"))?; - - // Use fused routing+fill kernel when conditions allow: - // - Not branching (branching learns order type via network heads) - // - Median vol > 0 (fill simulation needs vol context) - let use_routed = !self.hyperparams.use_branching && self.median_vol > 0.0; - - // Select actions via fused kernel — one launch, no intermediate GPU→CPU sync - let factored_tensor = if self.hyperparams.use_branching { - if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { - // Branching: 3-head epsilon-greedy → factored (0-44) directly - selector - .select_actions_branching(q_exp, q_ord, q_urg, epsilon) - .map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))? - } else { - // Branching Q forward failed — exposure-only → route on GPU - let exposure_tensor = selector - .select_actions(&batch_q_values, epsilon, batch_size, 5) - .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; - selector.route_exposure_to_factored( - &exposure_tensor, batch_size, - self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, - self.vol_ema as f32, self.median_vol as f32, - ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? - } - } else if use_routed { - // Fused: epsilon-greedy + routing + fill simulation in one kernel - let spread = self.hyperparams.avg_spread as f32; - let spread_bps = (self.hyperparams.avg_spread * 10000.0) as f32; - selector - .select_actions_routed( - &batch_q_values, epsilon, batch_size, 5, - batch_start as i32, - spread, spread, - self.vol_ema as f32, self.median_vol as f32, - spread_bps, 0.85, 0.30, 0.80, 0.50, 0.50, - ) - .map_err(|e| anyhow::anyhow!("GPU routed action selection failed: {e}"))? - } else { - // Exposure-only → route on GPU - let exposure_tensor = selector - .select_actions(&batch_q_values, epsilon, batch_size, 5) - .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; - selector.route_exposure_to_factored( - &exposure_tensor, batch_size, - self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, - self.vol_ema as f32, self.median_vol as f32, - ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? - }; - - // Per-element scalar readback (no to_vec1) - // narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar - let mut actions = Vec::with_capacity(batch_size); - for i in 0..batch_size { - let idx = factored_tensor.narrow(0, i, 1) - .and_then(|t| t.squeeze(0)) - .and_then(|t| t.to_scalar::()) - .map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?; - let action = FactoredAction::from_index(idx as usize) - .map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?; - actions.push(action); - } - Ok((actions, use_routed)) - } - - /// Epsilon-greedy action selection for single-step inference. - /// - /// This is batch_size=1 — the overhead of a CUDA kernel launch (~5us) exceeds - /// the benefit of fusing argmax+RNG for a single element. Candle's argmax + - /// to_scalar is already minimal for this path. Keep on CPU. - async fn epsilon_greedy_action(&self, state: &Tensor) -> Result { - use rand::Rng; - - let epsilon = self.get_epsilon().await? as f32; - let mut rng = rand::thread_rng(); - - if rng.gen::() < epsilon { - // Random action (exploration) over 5 exposure levels - Ok(rng.gen_range(0..5)) - } else { - // Greedy action (exploitation) - use actual Q-network - let agent = self.agent.read().await; - let q_values = agent.forward(state)?; - - // GPU-native argmax — single u32 scalar transfer instead of full Q-value vector - let best_action = q_values - .argmax(1) - .and_then(|t| t.squeeze(0)) - .and_then(|t| t.to_scalar::()) .map(|v| v as usize) - .ok() - .unwrap_or(2); // Default to HOLD (index 2) on error - - Ok(best_action) - } - } - - /// Calculate reward based on price movement - /// - /// # Arguments - /// * `current_close` - Current bar's close price - /// * `next_close` - Next bar's close price (target) - /// - /// # Returns - /// Normalized reward in [-1.0, 1.0] based on price change - fn calculate_reward(&self, current_close: f64, next_close: f64) -> f32 { - let price_change = next_close - current_close; - // Normalize by 10.0 for ES futures typical moves (±10 points) - // Clamp to [-1.0, 1.0] to prevent extreme rewards - (price_change / 10.0).clamp(-1.0, 1.0) as f32 - } - - /// Store experience in replay buffer - async fn store_experience(&self, experience: Experience) -> Result<()> { - let agent = self.agent.read().await; - agent - .store_experience(experience) - .map_err(|e| anyhow::anyhow!("Failed to store experience: {}", e))?; - Ok(()) - } - - /// Store a batch of experiences in one lock acquisition. - /// - /// Acquires the agent read lock once and delegates to - /// `DQNAgentType::store_experiences_batch` which holds the inner mutex - /// for the entire batch, reducing lock contention from O(n) to O(1). - async fn store_experiences_batch(&self, experiences: Vec) -> Result<()> { - if experiences.is_empty() { - return Ok(()); - } - let agent = self.agent.read().await; - agent - .store_experiences_batch(experiences) - .map_err(|e| anyhow::anyhow!("Failed to store experience batch: {}", e))?; - Ok(()) - } - - /// Check if we can train (buffer has enough samples) - async fn can_train(&self) -> Result { - let agent = self.agent.read().await; - Ok(agent.can_train()) - } - - /// Perform one training step using real DQN algorithm - /// - /// This method implements the core Deep Q-Learning algorithm: - /// 1. Sample batch from experience replay buffer - /// 2. Compute current Q-values: Q(s, a) - /// 3. Compute target Q-values: r + γ * max_a' Q_target(s', a') - /// 4. Calculate TD-error and MSE loss - /// 5. Backpropagate gradients and update Q-network - /// 6. Periodically update target network - /// - /// WAVE 26 P2.2: Now supports gradient accumulation for larger effective batch sizes - /// - /// Returns: (loss, avg_q_value, grad_norm) - async fn train_step(&mut self) -> Result<(f64, f64, f64)> { - let accumulation_steps = self.hyperparams.gradient_accumulation_steps; - - // OOM recovery loop: retry up to 3 times with halved batch size - const MAX_OOM_RETRIES: usize = 3; - - for retry in 0..=MAX_OOM_RETRIES { - let result = if accumulation_steps > 1 { - self.train_step_with_accumulation().await - } else { - self.train_step_single_batch().await - }; - - match result { - Ok(metrics) => return Ok(metrics), - Err(e) => { - // Check if this is an OOM error by inspecting the error chain - let err_str = format!("{:?}", e).to_lowercase(); - let is_oom = err_str.contains("out of memory") - || err_str.contains("oom") - || err_str.contains("cuda error 2") - || err_str.contains("cudamalloc") - || err_str.contains("failed to allocate"); - - if is_oom && retry < MAX_OOM_RETRIES { - let old_batch = self.current_batch_size; - self.current_batch_size = (old_batch / 2).max(1); - warn!( - "OOM detected (retry {}/{}): reducing batch size {} -> {}", - retry + 1, - MAX_OOM_RETRIES, - old_batch, - self.current_batch_size - ); - // Continue to next retry - } else { - return Err(e); - } - } - } - } - - Err(anyhow::anyhow!( - "Training failed after {} OOM retries", - MAX_OOM_RETRIES - )) - } - - /// Standard single-batch training step (no gradient accumulation) - /// - /// Pre-samples from PER buffer using READ lock before acquiring WRITE lock - /// for GPU training. This preserves PER IS-weights and indices for proper - /// importance sampling correction and priority updates. - async fn train_step_single_batch(&mut self) -> Result<(f64, f64, f64)> { - // Pre-sample batch OUTSIDE the write lock using read-only access to the buffer. - // PER IS-weights and indices are preserved for correct importance sampling. - // The write lock is only held during GPU forward/backward + optimizer step. - let explicit_batch = { - let agent = self.agent.read().await; - let buffer = agent.memory(); - let sample_size = self.current_batch_size; - buffer.can_sample(sample_size).then(|| { - buffer - .sample(sample_size) - .map_err(|e| anyhow::anyhow!("Failed to pre-sample batch: {}", e)) - }).transpose()? - }; // READ lock released here - - let mut agent = self.agent.write().await; - - // train_step returns GpuTrainResult with GPU-resident scalar tensors. - #[allow(unused_variables)] - let gpu_result = agent - .train_step(explicit_batch) - .map_err(|e| anyhow::anyhow!("Training step failed: {}", e))?; - - // GPU training guard: on-device NaN/loss-clip/grad-collapse checks. - // Zero cudaStreamSynchronize — kernel writes halt flags to pinned host memory. - #[cfg(feature = "cuda")] - let (loss_clipped, grad_norm) = { - // Lazy-init training guard on first call - if self.training_guard.is_none() && self.device.is_cuda() { - match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) { - Ok(guard) => { - info!("GPU training guard initialized"); - self.training_guard = Some(guard); - } - Err(e) => { - return Err(anyhow::anyhow!("GPU training guard init FAILED (no CPU fallback): {e}")); - } - } - } - - if let Some(ref mut guard) = self.training_guard { - let grad_collapse_threshold = - self.hyperparams.learning_rate as f32 - * self.hyperparams.gradient_collapse_multiplier as f32; - // Use original buffer_size (before AutoReplaySizer) for warmup guard - let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; - let past_warmup = self.gradient_logging_step as u64 > warmup_steps; - - let result = guard - .check_and_accumulate( - &gpu_result.loss_gpu, - &gpu_result.grad_norm_gpu, - 1e6_f32, // loss clip threshold - grad_collapse_threshold, - !past_warmup, - ) - .map_err(|e| anyhow::anyhow!("GPU guard check: {e}"))?; - - // Handle halt conditions - if result.halt_nan { - return Err(anyhow::anyhow!( - "NaN/Inf detected in loss ({}) or grad_norm ({})", - result.raw_loss, - result.raw_grad_norm - )); - } - if result.halt_loss_clip { - warn!( - "Loss clipped from {:.2e} to 1.0e6 (TD error explosion, epoch {})", - result.raw_loss, - self.loss_history.len() + 1 - ); - } - - // GPU guard path: collapse check only (no detect_dead_neurons GPU->CPU sync). - // Dead neuron detection runs at epoch boundary via log_diagnostics(). - agent.check_gradient_collapse(result.raw_grad_norm).map_err(|e| { - tracing::info!("Early stopping triggered (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - - (result.clipped_loss as f64, result.raw_grad_norm as f64) - } else { - return Err(anyhow::anyhow!( - "GPU training guard not initialized — CUDA device required for DQN training" - )); - } - }; - #[cfg(not(feature = "cuda"))] - return Err(anyhow::anyhow!( - "DQN training requires CUDA — enable the `cuda` feature" - )); - // Unreachable in non-cuda mode (return above), but Rust still name-checks. - #[cfg(not(feature = "cuda"))] - #[allow(unreachable_code)] - let (loss_clipped, grad_norm) = (0.0_f64, 0.0_f64); - - // Q-value estimation: periodic (every 50 steps) - self.q_estimation_counter += 1; - if self.q_estimation_counter % 50 == 1 { - // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) - #[cfg(feature = "cuda")] - let mut gpu_q_done = false; - #[cfg(feature = "cuda")] - { - if let Some(ref mut guard) = self.training_guard { - let buffer = agent.memory(); - if buffer.len() > 0 { - let sample_size = buffer.len().min(10); - let batch_sample = buffer - .sample(sample_size) - .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; - - let state_dim = agent.get_state_dim(); - let mut batch_tensor_opt: Option = None; - - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor_opt = Some( - gpu.states - .to_dtype(training_dtype(agent.device())) - .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, - ); - } - if batch_tensor_opt.is_none() { - let mut state_data = - Vec::with_capacity(sample_size * state_dim); - for exp in &batch_sample.experiences { - state_data.extend_from_slice(&exp.state); - } - if !state_data.is_empty() { - let tensor = Tensor::from_vec( state_data, - (sample_size, state_dim), - &self.device, - ) - .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; - batch_tensor_opt = Some(tensor); - } - } - - if let Some(ref batch_tensor) = batch_tensor_opt { - // Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync - agent.set_training_forward_active(true); - let batch_q_values = agent - .forward(batch_tensor) - .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; - agent.set_training_forward_active(false); - let num_actions = - batch_q_values.dims().get(1).copied().unwrap_or(5); - - // Divergence check on first sample - let first_q = batch_q_values - .i(0) - .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; - let div_result = guard - .qvalue_divergence(&first_q, num_actions, 10000.0) - .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; - agent - .log_q_values_from_stats( - div_result.q_min, - div_result.q_max, - div_result.q_mean, - div_result.q_variance, - num_actions, - ) - .map_err(|e| { - tracing::info!( - "Early stopping (Q-value divergence): {}", - e - ); - anyhow::anyhow!("Early stopping: {}", e) - })?; - - // Batch average via GPU reduction (one-step delay due to double-buffering) - let stats = guard - .qvalue_stats(&batch_q_values, sample_size, num_actions) - .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; - self.cached_avg_q = stats.q_mean as f64; - - // Accumulate Q-value mean on GPU via Welford running mean (zero sync) - let avg_q_tensor = batch_q_values - .max(1) - .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? - .mean_all() - .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; - guard - .accumulate_q_value(&avg_q_tensor) - .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; - - gpu_q_done = true; - } - } - } - } - // CUDA: GPU Q-value accumulation is mandatory — no CPU fallback. - #[cfg(feature = "cuda")] - if !gpu_q_done { - return Err(anyhow::anyhow!( - "GPU Q-value accumulation FAILED (no CPU fallback). \ - Check GpuTrainingGuard initialization." - )); - } - #[cfg(not(feature = "cuda"))] - { - self.cached_avg_q = - self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; - } - } - let avg_q_value = self.cached_avg_q; - - debug!("Gradient norm after clip (actual): {:.4}", grad_norm); - - self.gradient_logging_step += 1; - if self.gradient_logging_step % 10 == 0 { - debug!( - "Step {}: grad={:.4}, loss={:.4}", - self.gradient_logging_step, grad_norm, loss_clipped - ); - } - - Ok((loss_clipped, avg_q_value, grad_norm)) - } - - /// Training step with true gradient accumulation across N mini-batches. - /// - /// Unlike the previous implementation which ran N independent optimizer - /// steps, this version computes gradients for each mini-batch, accumulates - /// them, averages, and then applies a **single** optimizer step. This - /// simulates training with an effective batch size of - /// `accumulation_steps * batch_size` while keeping memory usage at - /// `batch_size`. - /// - /// Returns: (avg_loss, avg_q_value, final_grad_norm) - async fn train_step_with_accumulation(&mut self) -> Result<(f64, f64, f64)> { - let accumulation_steps = self.hyperparams.gradient_accumulation_steps; - - debug!( - "Starting true gradient accumulation with {} steps (effective batch: {})", - accumulation_steps, - self.current_batch_size * accumulation_steps - ); - - // Pre-sample ALL mini-batches using READ lock (no GPU contention). - let pre_sampled: Vec> = { - let agent = self.agent.read().await; - let buffer = agent.memory(); - let sample_size = self.current_batch_size; - let mut batches = Vec::with_capacity(accumulation_steps); - for step_idx in 0..accumulation_steps { - batches.push( - buffer.can_sample(sample_size).then(|| { - buffer.sample(sample_size).map_err(|e| { - anyhow::anyhow!( - "Failed to pre-sample batch (accum step {}): {}", - step_idx, - e - ) - }) - }).transpose()?, - ); - } - batches - }; // READ lock released - - let mut agent = self.agent.write().await; - - // === Phase 1: Accumulate gradients across N mini-batches === - let mut accumulated_grads: Option = None; - // Used by non-CUDA fallback and CUDA empty-tensor fallback paths. - #[allow(unused_mut, unused_assignments, unused_variables)] - let mut total_loss = 0.0_f64; - let mut all_td_errors = Vec::new(); - let mut all_indices = Vec::new(); - #[allow(unused_mut, unused_assignments, unused_variables)] - let mut final_grad_norm = 0.0_f32; - #[cfg(feature = "cuda")] - let mut gpu_td_errors: Vec = Vec::new(); - #[cfg(feature = "cuda")] - let mut gpu_indices: Vec = Vec::new(); - #[cfg(feature = "cuda")] - let mut gpu_loss_tensors: Vec = Vec::new(); - #[cfg(feature = "cuda")] - let mut gpu_grad_tensors: Vec = Vec::new(); - - for (step, batch) in pre_sampled.into_iter().enumerate() { - // Compute forward pass + backward WITHOUT optimizer step - let result = agent - .compute_gradients(batch) - .map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?; - - // Get vars for accumulation. Var is an Arc wrapper so cloning is cheap. - let vars: Vec = agent - .optimizer_vars() - .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?; - - crate::gradient_accumulation::accumulate_grads( - &mut accumulated_grads, - result.grads, - &vars, - ) - .map_err(|e| anyhow::anyhow!("Gradient accumulation step {} failed: {}", step, e))?; - - all_td_errors.extend(result.td_errors); - all_indices.extend(result.indices); - // GPU guard: check + accumulate loss/grad for this sub-step (borrows - // tensors before the move into gpu_*_tensors below). - #[cfg(feature = "cuda")] - { - if let (Some(ref loss_gpu), Some(ref gn_gpu)) = - (&result.loss_tensor_gpu, &result.grad_norm_gpu) - { - if let Some(ref mut guard) = self.training_guard { - let grad_collapse_threshold = - self.hyperparams.learning_rate as f32 - * self.hyperparams.gradient_collapse_multiplier as f32; - // Use original buffer_size (before AutoReplaySizer) for warmup guard - let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; - let past_warmup = self.gradient_logging_step as u64 > warmup_steps; - - let guard_result = guard.check_and_accumulate( - loss_gpu, - gn_gpu, - 1e6_f32, - grad_collapse_threshold, - !past_warmup, - ).map_err(|e| anyhow::anyhow!("GPU guard sub-step {}: {e}", step))?; - - if guard_result.halt_nan { - return Err(anyhow::anyhow!( - "NaN/Inf at accumulation sub-step {}: loss={}, grad={}", - step, guard_result.raw_loss, guard_result.raw_grad_norm - )); - } - if guard_result.halt_grad_collapse { - agent.check_gradient_collapse(guard_result.raw_grad_norm).map_err(|e| { - tracing::info!("Early stopping (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - } - } - } - } - #[cfg(feature = "cuda")] - { - if let Some(td_gpu) = result.td_errors_gpu { - gpu_td_errors.push(td_gpu); - } - if let Some(idx_gpu) = result.indices_gpu { - gpu_indices.push(idx_gpu); - } - if let Some(loss_gpu) = result.loss_tensor_gpu { - gpu_loss_tensors.push(loss_gpu); - } - if let Some(gn_gpu) = result.grad_norm_gpu { - gpu_grad_tensors.push(gn_gpu); - } - } - // CPU sentinel fallback (non-CUDA only) - #[cfg(not(feature = "cuda"))] - { - total_loss += result.loss as f64; - final_grad_norm = result.grad_norm; - } - } - - // === Phase 2: Average and apply gradients (single optimizer step) === - if let Some(ref mut grads) = accumulated_grads { - let vars: Vec = agent - .optimizer_vars() - .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?; - - crate::gradient_accumulation::scale_grads( - grads, - &vars, - 1.0 / accumulation_steps as f64, - ) - .map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?; - - #[cfg(feature = "cuda")] - let guard_active = self.training_guard.is_some(); - #[cfg(not(feature = "cuda"))] - let guard_active = false; - crate::gradient_accumulation::check_gradients_finite_guarded( - grads, - &vars, - guard_active, - ).map_err(|e| anyhow::anyhow!("Training halted: {}", e))?; - - agent - .apply_accumulated_gradients(grads) - .map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?; - } - - // === Phase 3: Bookkeeping === - #[cfg(feature = "cuda")] - { - // GPU PER path: concatenate GPU tensors and update in one shot - if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() { - let td_cat = candle_core::Tensor::cat(&gpu_td_errors, 0) - .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?; - let idx_cat = candle_core::Tensor::cat(&gpu_indices, 0) - .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?; - agent - .update_priorities_gpu(&idx_cat, &td_cat) - .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?; - } else if !all_indices.is_empty() { - agent - .update_priorities(&all_indices, &all_td_errors) - .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?; - } else { - // No priority updates needed (uniform buffer or empty batch) - } - } - #[cfg(not(feature = "cuda"))] - if !all_indices.is_empty() { - agent - .update_priorities(&all_indices, &all_td_errors) - .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?; - } - agent.step_replay_buffer(); - - // Single readback at accumulation boundary — prefer GPU guard accumulators - // (zero extra sync), fall back to cat+mean+to_scalar if guard absent. - #[cfg(feature = "cuda")] - let (avg_loss, final_grad_norm_f64) = { - if let Some(ref mut guard) = self.training_guard { - let (avg_l, avg_gn) = guard.read_accumulators() - .map_err(|e| anyhow::anyhow!("GPU guard read_accumulators: {e}"))?; - guard.reset_accumulators() - .map_err(|e| anyhow::anyhow!("GPU guard reset: {e}"))?; - (avg_l, avg_gn) - } else { - return Err(anyhow::anyhow!( - "GPU training guard not initialized — CUDA device required" - )); - } - }; - #[cfg(not(feature = "cuda"))] - let (avg_loss, final_grad_norm_f64) = ( - total_loss / accumulation_steps as f64, - final_grad_norm as f64, - ); - - // Gradient collapse detection (early stopping) -- no dead neuron GPU->CPU sync. - // Dead neuron detection runs at epoch boundary via log_diagnostics(). - agent - .check_gradient_collapse(final_grad_norm_f64 as f32) - .map_err(|e| { - tracing::info!("Early stopping triggered (gradient collapse): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - - // Clip averaged loss - let loss_clipped = if avg_loss > 1e6 { - warn!("Averaged loss clipped from {:.2e} to 1.0e6", avg_loss); - 1e6 - } else { - avg_loss - }; - - // Q-value estimation: periodic (every 50 steps) - self.q_estimation_counter += 1; - if self.q_estimation_counter % 50 == 1 { - // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) - #[cfg(feature = "cuda")] - let mut gpu_q_done = false; - #[cfg(feature = "cuda")] - { - if let Some(ref mut guard) = self.training_guard { - let buffer = agent.memory(); - if buffer.len() > 0 { - let sample_size = buffer.len().min(10); - let batch_sample = buffer - .sample(sample_size) - .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; - - let state_dim = agent.get_state_dim(); - let mut batch_tensor_opt: Option = None; - - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor_opt = Some( - gpu.states - .to_dtype(training_dtype(agent.device())) - .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, - ); - } - if batch_tensor_opt.is_none() { - let mut state_data = - Vec::with_capacity(sample_size * state_dim); - for exp in &batch_sample.experiences { - state_data.extend_from_slice(&exp.state); - } - if !state_data.is_empty() { - let tensor = Tensor::from_vec( state_data, - (sample_size, state_dim), - &self.device, - ) - .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? - .to_dtype(training_dtype(&self.device)) - .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; - batch_tensor_opt = Some(tensor); - } - } - - if let Some(ref batch_tensor) = batch_tensor_opt { - // Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync - agent.set_training_forward_active(true); - let batch_q_values = agent - .forward(batch_tensor) - .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; - agent.set_training_forward_active(false); - let num_actions = - batch_q_values.dims().get(1).copied().unwrap_or(5); - - // Divergence check on first sample - let first_q = batch_q_values - .i(0) - .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; - let div_result = guard - .qvalue_divergence(&first_q, num_actions, 10000.0) - .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; - agent - .log_q_values_from_stats( - div_result.q_min, - div_result.q_max, - div_result.q_mean, - div_result.q_variance, - num_actions, - ) - .map_err(|e| { - tracing::info!( - "Early stopping (Q-value divergence): {}", - e - ); - anyhow::anyhow!("Early stopping: {}", e) - })?; - - // Batch average via GPU reduction (one-step delay due to double-buffering) - let stats = guard - .qvalue_stats(&batch_q_values, sample_size, num_actions) - .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; - self.cached_avg_q = stats.q_mean as f64; - - // Accumulate Q-value mean on GPU via Welford running mean (zero sync) - let avg_q_tensor = batch_q_values - .max(1) - .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? - .mean_all() - .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; - guard - .accumulate_q_value(&avg_q_tensor) - .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; - - gpu_q_done = true; - } - } - } - } - // CUDA: GPU Q-value accumulation is mandatory — no CPU fallback. - #[cfg(feature = "cuda")] - if !gpu_q_done { - return Err(anyhow::anyhow!( - "GPU Q-value accumulation FAILED (no CPU fallback). \ - Check GpuTrainingGuard initialization." - )); - } - #[cfg(not(feature = "cuda"))] - { - self.cached_avg_q = - self.estimate_avg_q_value_with_early_stopping(&mut agent).await?; - } - } - let avg_q_value = self.cached_avg_q; - - Ok((loss_clipped, avg_q_value, final_grad_norm_f64)) - } - - /// Estimate average Q-value from replay buffer samples for monitoring - /// - /// WAVE 23 P0: Now includes Q-value divergence check (early stopping) - /// OPTIMIZATION: Batched Q-value estimation for 10× speedup via GPU parallelization - async fn estimate_avg_q_value_with_early_stopping(&self, agent: &mut DQNAgentType) -> Result { - // Get a few samples from the replay buffer to estimate Q-values - let buffer = agent.memory(); - - if buffer.len() == 0 { - return Ok(0.0); - } - - // Sample up to 10 experiences for Q-value estimation - let sample_size = buffer.len().min(10); - let batch_sample = buffer - .sample(sample_size) - .map_err(|e| anyhow::anyhow!("Failed to sample experiences: {}", e))?; - - let state_dim = agent.get_state_dim(); - - // GPU PER path: use gpu_batch.states directly (experiences vec is empty) - // CPU path: build tensor from experiences - #[allow(unused_assignments, unused_mut)] - let mut batch_tensor = None; - #[cfg(feature = "cuda")] - { - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor = Some( - gpu.states.to_dtype(training_dtype(agent.device())) - .map_err(|e| anyhow::anyhow!("GPU Q-est states dtype cast: {}", e))? - ); - } - } - let batch_tensor = if let Some(t) = batch_tensor { - t - } else { - let samples = &batch_sample.experiences; - let batched_states: Vec = samples.iter().flat_map(|exp| { - let mut s = exp.state.clone(); - s.resize(state_dim, 0.0); - s - }).collect(); - Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))? - .to_dtype(training_dtype(agent.device())) - .map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))? - }; - - // WAVE 23 P0 Fix: Check for Q-value divergence (early stopping) - // This calls log_q_values() which returns Err if divergence detected for consecutive checks - agent.log_q_values(&batch_tensor) - .map_err(|e| { - tracing::info!("🛑 Early stopping triggered (Q-value divergence): {}", e); - anyhow::anyhow!("Early stopping: {}", e) - })?; - - // Single forward pass for all samples (10× faster than sequential) - let batch_q_values = agent - .forward(&batch_tensor) - .map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?; - - // Get max Q-value per sample across action dimension - let max_q_values = batch_q_values - .max(1) - .map_err(|e| anyhow::anyhow!("Failed to compute max Q-values: {}", e))?; - - // Compute average across batch - let avg_q = max_q_values - .mean_all() - .map_err(|e| anyhow::anyhow!("Failed to compute mean Q-value: {}", e))? - .to_scalar::() .map_err(|e| anyhow::anyhow!("Failed to extract average Q-value: {}", e))? - as f64; - - Ok(avg_q) - } - - /// Epoch-end Q-value diagnostics: gap analysis + per-action averages. - /// - /// Merges the former `compute_q_gap_for_epoch` and `compute_per_action_q_values` - /// into a single forward pass + readback, eliminating one redundant buffer sample, - /// forward pass, and `to_vec2` GPU-CPU transfer per epoch. - /// - /// Returns (gap_stats, per_action_avgs) where: - /// - gap_stats: (mean_gap, min_gap, max_gap) of Q_best - Q_second_best - /// - per_action_avgs: `[f64; 5]` averages (one per exposure action) - async fn compute_epoch_q_diagnostics(&self) -> Option<( - (f64, f64, f64), - [f64; 5], - )> { - let agent = self.agent.read().await; - let buffer = agent.memory(); - - if buffer.len() < 10 { - return None; - } - - // Use the larger sample size (200) for both diagnostics - let sample_size = buffer.len().min(200); - let batch_sample = match buffer.sample(sample_size) { - Ok(s) => s, - Err(_) => return None, - }; - - let state_dim = agent.get_state_dim(); - - // GPU PER path: use gpu_batch.states directly (experiences vec is empty) - #[allow(unused_mut)] - let mut batch_tensor_opt: Option = None; - #[cfg(feature = "cuda")] - { - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor_opt = gpu.states - .to_dtype(training_dtype(agent.device())) - .ok(); - } - } - let batch_tensor = if let Some(t) = batch_tensor_opt { - t - } else { - let batched_states: Vec = batch_sample - .experiences - .iter() - .flat_map(|exp| { - let mut s = exp.state.clone(); - s.resize(state_dim, 0.0); - s - }) - .collect(); - let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) { Ok(t) => t, - Err(_) => return None, - }; - match t.to_dtype(training_dtype(agent.device())) { - Ok(t) => t, - Err(_) => return None, - } - }; - - // Single forward pass for both gap and per-action diagnostics - let batch_q_values = match agent.forward(&batch_tensor) { - Ok(q) => q, - Err(_) => return None, - }; - - // All paths use GPU-native diagnostics — no to_vec2 readback - #[cfg(feature = "cuda")] - if self.device.is_cuda() { - return compute_q_diagnostics_gpu(&batch_q_values).ok(); - } - - // Non-CUDA: compute diagnostics via tensor ops (no to_vec2) - // Sort Q-values descending per row, gap = sorted[0] - sorted[1] - let sorted = match batch_q_values.sort_last_dim(false) { - Ok((s, _)) => s, - Err(_) => return None, - }; - let n_actions = batch_q_values.dims().get(1).copied().unwrap_or(5); - if n_actions < 2 { return None; } - - let best = match sorted.narrow(1, 0, 1) { - Ok(t) => t.flatten_all().unwrap_or(sorted.clone()), - Err(_) => return None, - }; - let second_best = match sorted.narrow(1, 1, 1) { - Ok(t) => t.flatten_all().unwrap_or(sorted.clone()), - Err(_) => return None, - }; - let gaps_tensor = match best.sub(&second_best) { - Ok(t) => t, - Err(_) => return None, - }; - - let gap_mean = gaps_tensor.mean_all().ok()?.to_scalar::().ok()? as f64; - let gap_min = gaps_tensor.min(0).ok()?.to_scalar::().ok()? as f64; - let gap_max = gaps_tensor.max(0).ok()?.to_scalar::().ok()? as f64; - - // Per-action average Q-values via mean(dim=0) - let per_action = match batch_q_values.mean(0) { - Ok(t) => t, - Err(_) => return None, - }; - - let mut avgs = [0.0_f64; 5]; - for i in 0..5_usize.min(n_actions) { - if let Ok(v) = per_action.narrow(0, i, 1).and_then(|t| t.to_scalar::()) { - avgs[i] = v as f64; - } - } - - Some(((gap_mean, gap_min, gap_max), avgs)) - } - - /// Get current epsilon value - async fn get_epsilon(&self) -> Result { - let agent = self.agent.read().await; - Ok(agent.get_epsilon() as f64) - } - - /// Set epsilon value (used for deterministic evaluation) - async fn set_epsilon(&self, epsilon: f64) -> Result<()> { - let mut agent = self.agent.write().await; - agent.set_epsilon(epsilon); - Ok(()) - } - - /// Get best validation loss achieved during training - /// - /// Returns the lowest validation loss seen across all epochs. - /// Used by hyperopt adapter to optimize for generalization. - pub fn get_best_val_loss(&self) -> f64 { - self.best_val_loss - } - - /// Get epoch number where best validation loss was achieved - /// - /// Returns the 1-indexed epoch number with the best validation loss. - pub fn get_best_epoch(&self) -> usize { - self.best_epoch - } - - /// Get validation data for backtest integration - /// - /// Returns a reference to the validation dataset for hyperopt backtest evaluation. - /// Each entry contains a FeatureVector (42 market + 3 portfolio = 45 dims) and the corresponding target values. - /// Used by hyperopt adapter to run backtests on unseen data after training. - pub fn get_val_data(&self) -> &[(FeatureVector, Vec)] { - &self.val_data - } - - /// Convert feature vector to state tensor for action selection - /// - /// Public wrapper around internal state conversion for hyperopt backtest integration. - /// Converts a 42-dimensional feature vector to a 45-dimensional state tensor - /// suitable for DQN agent's select_action method. - /// - /// # Arguments - /// - /// * `feature_vec` - 42-dimensional market feature vector - /// * `close_price` - Current close price for portfolio feature calculation - /// - /// # Returns - /// - /// Result containing the 45-dimensional state tensor ready for model inference. - /// Portfolio features (last 3 dimensions) are populated via PortfolioTracker. - pub fn convert_to_state( - &self, - feature_vec: &FeatureVector, - close_price: f64, - ) -> Result { - let close = rust_decimal::Decimal::try_from(close_price) - .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; - - // Use internal conversion method (returns TradingState) - let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?; - - // Convert TradingState to flat vector, pad for tensor core alignment - let state_vec = trading_state.to_vector(); - let raw_dim = state_vec.len(); - let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); - let padded: Vec = if aligned > raw_dim { - let mut v = state_vec.to_vec(); - v.resize(aligned, 0.0); - v - } else { - state_vec.to_vec() - }; - - // Convert to Tensor using trainer's device (GPU or CPU) - Tensor::new(padded.as_slice(), &self.device) - .context("Failed to create state tensor from TradingState") - } - - /// Convert feature vector to flat state Vec (CPU only, no GPU tensor). - /// - /// Same as `convert_to_state` but returns the raw vector instead of a GPU tensor. - /// Used by chunked batch inference to avoid per-bar GPU allocations. - pub fn convert_to_state_vec( - &self, - feature_vec: &FeatureVector, - close_price: f64, - ) -> Result> { - let close = rust_decimal::Decimal::try_from(close_price) - .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; - let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?; - Ok(trading_state.to_vector()) - } - - /// Convert feature vector to flat state Vec with OFI features at the given index. - /// - /// Same as `convert_to_state_vec` but injects OFI features from the preloaded - /// array at `ofi_index`, preventing train/eval feature mismatch. - pub fn convert_to_state_vec_with_ofi( - &self, - feature_vec: &FeatureVector, - close_price: f64, - ofi_index: usize, - ) -> Result> { - let close = rust_decimal::Decimal::try_from(close_price) - .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; - let trading_state = self.feature_vector_to_state_with_ofi(feature_vec, Some(close), Some(ofi_index))?; - Ok(trading_state.to_vector()) - } - - /// Update portfolio tracker to reflect current position from backtest engine. - /// - /// Called between chunks so the next chunk's portfolio features - /// accurately reflect the current position (direction, value, exposure). - pub fn set_portfolio_for_backtest( - &mut self, - position_size: f32, - entry_price: f32, - current_price: f32, - ) { - self.portfolio_tracker = PortfolioTracker::new( - self.portfolio_tracker.initial_capital(), - self.portfolio_tracker.spread(), - 0.0, - ); - if position_size.abs() > f32::EPSILON { - self.portfolio_tracker - .set_position_direct(position_size, entry_price, current_price); - } - } - - /// Get the device used by this trainer - pub fn device(&self) -> &candle_core::Device { - &self.device - } - - /// Get access to the DQN agent - /// - /// Returns a reference to the Arc> for checkpoint saving. - /// Used by hyperopt adapter to save model weights after training. - pub fn get_agent(&self) -> &Arc> { - &self.agent - } - - /// Get reference to training hyperparameters - /// - /// Returns a reference to the DQN hyperparameters used for this trainer. - /// Used by tests to validate configuration. - pub fn hyperparams(&self) -> &DQNHyperparameters { - &self.hyperparams - } - - /// Get current learning rate from scheduler - /// - /// Returns the current learning rate after applying warmup and decay. - /// Used by tests and monitoring to track LR schedule. - pub fn get_current_lr(&self) -> f64 { - self.lr_scheduler.get_lr() - } - - /// Serialize model to bytes with architecture metadata embedded in safetensors header. - /// - /// For RegimeConditional agents, serializes ALL 3 heads (trending, ranging, - /// volatile) into a single safetensors file using prefixed tensor names - /// (`trending__`, `ranging__`, `volatile__`). This ensures walk-forward - /// checkpoint restore loads all heads, not just the trending head. - pub async fn serialize_model(&self) -> Result> { - let agent = self.agent.read().await; - - let tensors: std::collections::HashMap = match &*agent { - crate::trainers::dqn::DQNAgentType::RegimeConditional(regime) => { - let mut all_tensors = std::collections::HashMap::new(); - for (prefix, head_opt) in [ - ("trending__", regime.get_trending_head()), - ("ranging__", regime.get_ranging_head()), - ("volatile__", regime.get_volatile_head()), - ] { - let head = head_opt.ok_or_else(|| { - anyhow::anyhow!("Missing {} head for serialization", prefix) - })?; - let vars = head.get_q_network_vars(); - let vars_data = vars.data().lock().map_err(|_| { - anyhow::anyhow!("Failed to lock VarMap for {} head", prefix) - })?; - for (name, var) in vars_data.iter() { - all_tensors.insert( - format!("{}{}", prefix, name), - var.as_tensor().clone(), - ); - } - } - all_tensors - } - _ => { - let vars = agent.get_q_network_vars(); - let vars_data = vars.data().lock().map_err(|_| { - anyhow::anyhow!("Failed to lock VarMap for serialization") - })?; - vars_data - .iter() - .map(|(name, var)| (name.clone(), var.as_tensor().clone())) - .collect() - } - }; - - // Embed architecture metadata in safetensors header - let arch_metadata = Some(agent.checkpoint_metadata()); - let data = safetensors::serialize(&tensors, &arch_metadata) - .map_err(|e| anyhow::anyhow!("Failed to serialize safetensors: {}", e))?; - - Ok(data) - } - - /// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`). - /// - /// The trainer's `train_epoch` lazily uploads data on first call. - /// Use this to provide data that was uploaded in advance by a - /// `DoubleBufferedLoader`, skipping the per-fold upload latency. - pub fn set_gpu_data(&mut self, data: DqnGpuData) { - info!( - "DqnTrainer: injected pre-uploaded GPU data ({} bars, {:.1} MB)", - data.num_bars, - data.vram_bytes() as f64 / 1_048_576.0, - ); - self.gpu_data = Some(data); - } - - /// Drop cached GPU data, freeing VRAM for the next fold. - pub fn clear_gpu_data(&mut self) { - if self.gpu_data.is_some() { - info!("DqnTrainer: cleared GPU data (VRAM freed)"); - self.gpu_data = None; - } - } - - /// BUG #38 FIX: Clear replay buffer of contaminated experiences - pub async fn clear_replay_buffer(&mut self) -> Result<()> { - let mut agent = self.agent.write().await; - agent.clear_replay_buffer().map_err(|e| { - anyhow::anyhow!("Failed to clear replay buffer: {}", e) - })?; - let buffer_size = agent.get_replay_buffer_size().unwrap_or(0); - info!("Replay buffer cleared successfully. Current size: {}", buffer_size); - Ok(()) - } - - /// BUG #38 FIX: Reset target network to match current network - pub async fn reset_target_network(&mut self) -> Result<()> { - let mut agent = self.agent.write().await; - agent.reset_target_network().map_err(|e| { - anyhow::anyhow!("Failed to reset target network: {}", e) - })?; - info!("Target network reset successfully"); - Ok(()) - } - - - - /// Get current training metrics - pub async fn get_metrics(&self) -> TrainingMetrics { - self.metrics.read().await.clone() - } - - /// Get per-epoch training loss history (for smoke test verification) - pub fn loss_history(&self) -> &[f64] { - &self.loss_history - } - - /// Get per-epoch validation loss history - pub fn val_loss_history(&self) -> &[f64] { - &self.val_loss_history - } - - /// Get current epsilon from the DQN agent - pub async fn get_agent_epsilon(&self) -> f32 { - let agent_lock = self.agent.read().await; - agent_lock.get_epsilon() - } -} - -// --------------------------------------------------------------------------- -// GPU Q-value diagnostics (Task 6) -// --------------------------------------------------------------------------- - -/// Compute Q-value gap and per-action averages on GPU. -/// Returns (mean_gap, min_gap, max_gap, per_action_avgs[5]). -/// Single 8-float readback at epoch end. -#[cfg(feature = "cuda")] -fn compute_q_diagnostics_gpu( - q_values: &Tensor, // [batch, 5] -) -> candle_core::Result<((f64, f64, f64), [f64; 5])> { - // sort_last_dim returns (sorted_values, indices) — destructure the tuple - let (sorted, _indices) = q_values.sort_last_dim(true)?; // descending - let best = sorted.narrow(1, 0, 1)?; - let second = sorted.narrow(1, 1, 1)?; - let gaps = best.sub(&second)?; - - // Batch all gap stats into a single tensor to minimize readbacks: - // Note: In candle-core (git 671de1d), min(D)/max(D) return Result, - // NOT Result<(Tensor, Tensor)>. Flatten first for scalar reduction. - let gaps_flat = gaps.flatten_all()?; - let mean_gap = gaps_flat.mean_all()?; // scalar tensor - let min_gap = gaps_flat.min(0)?; // scalar tensor - let max_gap = gaps_flat.max(0)?; // scalar tensor - let _gap_stats = Tensor::cat( - &[&mean_gap.unsqueeze(0)?, &min_gap.unsqueeze(0)?, &max_gap.unsqueeze(0)?], 0 - )?; - - // Per-action means: mean along batch dim [5] - let per_action = q_values.mean(0)?; - - // 8 scalar readbacks: 3 gap stats + 5 per-action means (no to_vec1) - let mean_g = mean_gap.to_scalar::()? as f64; - let min_g = min_gap.to_scalar::()? as f64; - let max_g = max_gap.to_scalar::()? as f64; - let mut avgs = [0.0_f64; 5]; - for i in 0..5_usize { - if i < per_action.elem_count() { - avgs[i] = per_action.narrow(0, i, 1)? - .to_scalar::()? as f64; - } - } - - Ok(((mean_g, min_g, max_g), avgs)) -} - -// --------------------------------------------------------------------------- -// GPU batch → Experience conversion (Phase 3) -// --------------------------------------------------------------------------- - -/// Convert a GPU `ExperienceBatch` (flat arrays from CUDA kernel) into `Vec` -/// for insertion into the DQN replay buffer. -/// -/// `state_dim` must be the **aligned** dimension (e.g. 56 with OFI, 48 without) matching -/// the value injected as `STATE_DIM` into the CUDA kernel. The GPU outputs states at this -/// stride, so using the raw (unaligned) dim would corrupt all samples after the first. -#[cfg(feature = "cuda")] -fn gpu_batch_to_experiences( - batch: &crate::cuda_pipeline::gpu_experience_collector::ExperienceBatch, - state_dim: usize, -) -> Vec { - let total = batch.n_episodes * batch.timesteps; - let mut experiences = Vec::with_capacity(total); - - for ep in 0..batch.n_episodes { - for t in 0..batch.timesteps { - let idx = ep * batch.timesteps + t; - let state_start = idx * state_dim; - let state_end = state_start + state_dim; - // Skip out-of-bounds entries early instead of creating empty vecs - let state_slice = match batch.states.get(state_start..state_end) { - Some(s) if s.len() == state_dim => s, - _ => continue, - }; - - // Clamp action to valid u8 range (0..4 for 5 exposure levels) - let action = batch.actions.get(idx).copied().unwrap_or(0).clamp(0, 4) as u8; - let reward = batch.rewards.get(idx).copied().unwrap_or(0.0); - let done = batch.done_flags.get(idx).copied().unwrap_or(0) != 0; - - // next_state = states[t+1] within same episode, or current state if done/last - let next_state = if !done && t + 1 < batch.timesteps { - let next_idx = ep * batch.timesteps + t + 1; - let ns_start = next_idx * state_dim; - let ns_end = ns_start + state_dim; - match batch.states.get(ns_start..ns_end) { - Some(s) if s.len() == state_dim => s.to_vec(), - _ => state_slice.to_vec(), // fallback to current state - } - } else { - state_slice.to_vec() - }; - - experiences.push(Experience::new(state_slice.to_vec(), action, reward, next_state, done)); - } - } - experiences -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::hyperopt::ParameterSpace; - use std::sync::OnceLock; - - /// Shared CUDA device across all trainer tests. - /// - /// Root cause fix for CUBLAS_STATUS_NOT_INITIALIZED cascades: each - /// `Device::cuda_if_available(0)` creates a new cuBLAS handle. With - /// 400+ tests doing this in rapid succession (even with --test-threads=1), - /// the driver's internal handle pool is exhausted. Sharing one device - /// eliminates the churn entirely. - static SHARED_DEVICE: OnceLock = OnceLock::new(); - - fn shared_cuda_device() -> Device { - // Initialize tracing so kernel compilation/launch logs are visible. - static TRACING_INIT: std::sync::Once = std::sync::Once::new(); - TRACING_INIT.call_once(|| { - let filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); - tracing_subscriber::fmt() - .with_env_filter(filter) - .with_test_writer() - .try_init() - .ok(); - }); - SHARED_DEVICE - .get_or_init(|| { - Device::cuda_if_available(0).unwrap_or(Device::Cpu) - }) - .clone() - } - - // Helper function to create test hyperparameters - // Uses conservative defaults suitable for testing - fn create_test_params() -> DQNHyperparameters { - let mut params = DQNHyperparameters::conservative(); - // Production default: branching DQN (3 heads: exposure, order, urgency). - // Always enabled — the warp-cooperative kernel on H100 requires it. - params.use_branching = true; - params.hidden_dim_base = Some(32); // Small for fast test iterations - params.buffer_size = 10_000; - params - } - - fn create_test_trainer() -> Result { - DQNTrainer::new_with_device(create_test_params(), shared_cuda_device()) - } - - fn create_test_trainer_with(params: DQNHyperparameters) -> Result { - DQNTrainer::new_with_device(params, shared_cuda_device()) - } - - /// Pad a TradingState's regime_features so that `state.dimension()` matches the - /// trainer's aligned state_dim (e.g. 45→48 on CUDA due to tensor core alignment). - fn pad_state_to_aligned(state: &mut TradingState, trainer: &DQNTrainer) { - let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores( - state.dimension(), - &trainer.device, - ); - let pad = aligned_dim.saturating_sub(state.dimension()); - if pad > 0 { - state.regime_features.extend(vec![0.0_f32; pad]); - } - } - - #[tokio::test] - async fn test_dqn_trainer_creation() { - let hyperparams = create_test_params(); - let trainer = create_test_trainer_with(hyperparams); - - assert!( - trainer.is_ok(), - "Failed to create DQN trainer: {:?}", - trainer.err() - ); - } - - #[tokio::test] - async fn test_batch_size_validation() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 500; - - // VRAM ceiling clamps if needed, never rejects - let trainer = create_test_trainer_with(hyperparams); - assert!( - trainer.is_ok(), - "Should clamp oversized batch, not reject: {:?}", - trainer.err() - ); - } - - #[tokio::test] - async fn test_feature_vector_to_state() { - let hyperparams = create_test_params(); - let trainer = create_test_trainer_with(hyperparams).unwrap(); - - // Create a synthetic 42-dim feature vector (42 market features) - let mut feature_vec = [0.0; 42]; - feature_vec[0] = 4000.0; // open - feature_vec[1] = 4010.0; // high - feature_vec[2] = 3990.0; // low - feature_vec[3] = 4005.0; // close - feature_vec[4] = 1000.0; // volume - // Fill remaining features with synthetic data - for i in 5..42 { - feature_vec[i] = (i as f64) * 0.1; - } - - let close_price = - rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); - let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)); - - assert!( - state.is_ok(), - "Failed to convert feature vector: {:?}", - state.err() - ); - - let state = state.unwrap(); - // State dimension is 45 (42 market + 3 portfolio) - // - Market features: 0-41 (42 features) - // - Portfolio features: 42-44 (3 features, populated by PortfolioTracker) - assert_eq!( - state.dimension(), - 45, - "State dimension should be 45 (42 market + 3 portfolio features)" - ); - } - - #[cfg_attr(not(feature = "cuda"), ignore)] - #[tokio::test] - async fn test_batched_action_selection() { - let hyperparams = create_test_params(); - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - // Create multiple synthetic states for batched action selection - let batch_size = 10; - let mut states = Vec::with_capacity(batch_size); - - for i in 0..batch_size { - let mut feature_vec = [0.0; 42]; // 42 market features - // Create varied states for testing - feature_vec[0] = 4000.0 + (i as f64 * 10.0); // open - feature_vec[1] = 4010.0 + (i as f64 * 10.0); // high - feature_vec[2] = 3990.0 + (i as f64 * 10.0); // low - feature_vec[3] = 4005.0 + (i as f64 * 10.0); // close - feature_vec[4] = 1000.0 + (i as f64 * 100.0); // volume - - // Fill remaining features - for j in 5..42 { - feature_vec[j] = (j as f64 + i as f64) * 0.1; - } - - let close_price = rust_decimal::Decimal::try_from(feature_vec[3]) - .unwrap_or(rust_decimal::Decimal::ZERO); - let mut state = trainer - .feature_vector_to_state(&feature_vec, Some(close_price)) - .unwrap(); - pad_state_to_aligned(&mut state, &trainer); - states.push(state); - } - - // Test batched action selection - let actions_result = trainer.select_actions_batch(&states).await; - - assert!( - actions_result.is_ok(), - "Batched action selection failed: {:?}", - actions_result.err() - ); - - let actions = actions_result.unwrap(); - assert_eq!( - actions.len(), - batch_size, - "Expected {} actions, got {}", - batch_size, - actions.len() - ); - - // Verify all actions have valid exposure indices (0-4) - for (i, action) in actions.iter().enumerate() { - let exp_idx = action.exposure as usize; - assert!( - exp_idx < 5, - "Action {} has invalid exposure index {}: {:?}", - i, - exp_idx, - action - ); - } - } - - #[cfg_attr(not(feature = "cuda"), ignore)] - #[tokio::test] - async fn test_batched_vs_sequential_action_selection_consistency() { - let hyperparams = create_test_params(); - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - // Create test states - let batch_size = 5; - let mut states = Vec::with_capacity(batch_size); - - for i in 0..batch_size { - let mut feature_vec = [0.0; 42]; // 42 market features - feature_vec[0] = 4000.0 + (i as f64 * 50.0); - feature_vec[1] = 4050.0 + (i as f64 * 50.0); - feature_vec[2] = 3950.0 + (i as f64 * 50.0); - feature_vec[3] = 4025.0 + (i as f64 * 50.0); - feature_vec[4] = 5000.0 + (i as f64 * 500.0); - - for j in 5..42 { - feature_vec[j] = (j as f64) * 0.5 + (i as f64); - } - - let close_price = rust_decimal::Decimal::try_from(feature_vec[3]) - .unwrap_or(rust_decimal::Decimal::ZERO); - let mut state = trainer - .feature_vector_to_state(&feature_vec, Some(close_price)) - .unwrap(); - pad_state_to_aligned(&mut state, &trainer); - states.push(state); - } - - // Get batched actions (GPU-optimized) - let batched_actions = trainer.select_actions_batch(&states).await.unwrap(); - - // Both should return valid actions - assert_eq!( - batched_actions.len(), - batch_size, - "Batched action count mismatch" - ); - - // Verify all actions have valid exposure indices (0-4) - for action in &batched_actions { - let exp_idx = action.exposure as usize; - assert!(exp_idx < 5, "Invalid exposure index {}: {:?}", exp_idx, action); - } - } - - #[tokio::test] - async fn test_empty_batch_handling() { - let hyperparams = create_test_params(); - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - let empty_states: Vec = Vec::new(); - let result = trainer.select_actions_batch(&empty_states).await; - - assert!(result.is_ok(), "Empty batch should be handled gracefully"); - assert_eq!( - result.unwrap().len(), - 0, - "Empty batch should return empty actions" - ); - } - - #[tokio::test] - async fn test_zero_batch_size_handling() { - // Test DQN rejects zero batch size - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 0; - - let result = create_test_trainer_with(hyperparams); - - // Should fail with descriptive error - assert!( - result.is_err(), - "DQN should reject zero batch size, but got: {:?}", - result - ); - - // Error message should mention batch size - let error_msg = result.unwrap_err().to_string(); - assert!( - error_msg.to_lowercase().contains("batch"), - "Error message should mention batch size, got: {}", - error_msg - ); - } - - // ===== Agent 23 Test #6: Batch Size Mismatch Validation Tests ===== - - /// Production-critical test: Verify trainer handles batch smaller than configured - #[cfg_attr(not(feature = "cuda"), ignore)] - #[tokio::test] - async fn test_batch_size_mismatch_smaller_than_configured() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 32; - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - // Create batch with 16 states (half of configured 32) - let mut feature_vec = [0.0; 42]; // 42 market features - for i in 0..4 { - feature_vec[i] = 4000.0 + (i as f64 * 10.0); - } - for i in 5..42 { - feature_vec[i] = (i as f64) * 0.1; - } - - let close_price = - rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); - let mut state = trainer - .feature_vector_to_state(&feature_vec, Some(close_price)) - .unwrap(); - pad_state_to_aligned(&mut state, &trainer); - let smaller_batch = vec![state.clone(); 16]; - - let result = trainer.select_actions_batch(&smaller_batch).await; - assert!( - result.is_ok(), - "DQN should handle smaller batches: {:?}", - result.err() - ); - assert_eq!( - result.unwrap().len(), - 16, - "Should return action for each state" - ); - } - - /// Production-critical test: Verify trainer handles batch larger than configured - #[cfg_attr(not(feature = "cuda"), ignore)] - #[tokio::test] - async fn test_batch_size_mismatch_larger_than_configured() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 16; - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - // Create batch with 64 states (4x configured 16) - let mut feature_vec = [0.0; 42]; // 42 market features - for i in 0..4 { - feature_vec[i] = 4000.0 + (i as f64 * 10.0); - } - for i in 5..42 { - feature_vec[i] = (i as f64) * 0.1; - } - - let close_price = - rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); - let mut state = trainer - .feature_vector_to_state(&feature_vec, Some(close_price)) - .unwrap(); - pad_state_to_aligned(&mut state, &trainer); - let larger_batch = vec![state.clone(); 64]; - - let result = trainer.select_actions_batch(&larger_batch).await; - assert!( - result.is_ok(), - "DQN should handle larger batches: {:?}", - result.err() - ); - assert_eq!( - result.unwrap().len(), - 64, - "Should return action for each state" - ); - } - - /// Production-critical test: Verify empty batch handling - #[tokio::test] - async fn test_empty_batch_returns_empty_actions() { - let mut trainer = create_test_trainer().unwrap(); - let empty_batch: Vec = vec![]; - - let result = trainer.select_actions_batch(&empty_batch).await; - assert!(result.is_ok(), "Should handle empty batch gracefully"); - assert_eq!( - result.unwrap().len(), - 0, - "Empty batch should return empty actions" - ); - } - - /// Production-critical test: Verify single-sample batch handling - #[cfg_attr(not(feature = "cuda"), ignore)] - #[tokio::test] - async fn test_single_sample_batch() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 32; - let mut trainer = create_test_trainer_with(hyperparams).unwrap(); - - let mut feature_vec = [0.0; 42]; // 42 market features - for i in 0..4 { - feature_vec[i] = 4000.0; - } - for i in 5..42 { - feature_vec[i] = (i as f64) * 0.1; - } - - let close_price = - rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); - let mut state = trainer - .feature_vector_to_state(&feature_vec, Some(close_price)) - .unwrap(); - pad_state_to_aligned(&mut state, &trainer); - let single_batch = vec![state]; - - let result = trainer.select_actions_batch(&single_batch).await; - assert!( - result.is_ok(), - "Should handle single-sample batch: {:?}", - result.err() - ); - assert_eq!(result.unwrap().len(), 1, "Should return exactly one action"); - } - - /// Large batch sizes are accepted (VRAM ceiling is the only cap) - #[test] - fn test_large_batch_size_accepted() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 2048; - - let result = create_test_trainer_with(hyperparams); - assert!( - result.is_ok(), - "Should accept large batch sizes within VRAM ceiling: {:?}", - result.err() - ); - } - - /// Production-critical test: Non-power-of-2 batch sizes - #[tokio::test] - async fn test_non_power_of_two_batch_size() { - let mut hyperparams = create_test_params(); - hyperparams.batch_size = 13; // Not a power of 2 - - let result = create_test_trainer_with(hyperparams); - assert!( - result.is_ok(), - "Should accept non-power-of-2 batch sizes: {:?}", - result.err() - ); - } - - /// Production-critical test: Train with empty dataset doesn't crash - #[tokio::test] - async fn test_train_with_empty_data_completes_gracefully() { - let mut params = create_test_params(); - params.epochs = 5; // Short run — just checking it doesn't panic - params.early_stopping_enabled = false; - params.gradient_collapse_patience = 1000; - params.buffer_size = 1000; - let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); - let mut trainer = DQNTrainer::new_with_device(params, device).unwrap(); - let empty_data: Vec<(FeatureVector, Vec)> = vec![]; - let checkpoint_callback = |_, _, _| Ok(String::new()); - - let result = trainer - .train_with_data_full_loop(&empty_data, checkpoint_callback) - .await; - - assert!( - result.is_err(), - "Training with empty data should return an error (no CPU fallback)" - ); - } - - /// Test reward function calculates actual price changes correctly - #[test] - fn test_reward_function_price_changes() { - let trainer = create_test_trainer().unwrap(); - - // Test upward price move (+14.25 points, should clamp to +1.0) - let reward_up = trainer.calculate_reward(5900.0, 5914.25); - assert!( - (reward_up - 1.0).abs() < 1e-6, - "Upward move should return +1.0 (clamped), got: {}", - reward_up - ); - - // Test downward price move (-14.25 points, should clamp to -1.0) - let reward_down = trainer.calculate_reward(5914.25, 5900.0); - assert!( - (reward_down - (-1.0)).abs() < 1e-6, - "Downward move should return -1.0 (clamped), got: {}", - reward_down - ); - - // Test flat market (0 points, should return 0.0) - let reward_flat = trainer.calculate_reward(5900.0, 5900.0); - assert!( - reward_flat.abs() < 1e-6, - "Flat market should return 0.0, got: {}", - reward_flat - ); - - // Test small upward move (+5 points, should return +0.5) - let reward_small_up = trainer.calculate_reward(5900.0, 5905.0); - assert!( - (reward_small_up - 0.5).abs() < 1e-6, - "Small upward move (+5) should return +0.5, got: {}", - reward_small_up - ); - - // Test small downward move (-5 points, should return -0.5) - let reward_small_down = trainer.calculate_reward(5905.0, 5900.0); - assert!( - (reward_small_down - (-0.5)).abs() < 1e-6, - "Small downward move (-5) should return -0.5, got: {}", - reward_small_down - ); - - // Test unclamped move (+3 points, should return +0.3) - let reward_unclamped = trainer.calculate_reward(5900.0, 5903.0); - assert!( - (reward_unclamped - 0.3).abs() < 1e-6, - "Move of +3 points should return +0.3, got: {}", - reward_unclamped - ); - } - - #[test] - fn test_dynamic_batch_size_l4() { - // L4 has 24GB VRAM — HardwareBudget should allow batch_size >> 230 - let budget = crate::hyperopt::HardwareBudget { - gpu_memory_mb: 24_000, - gpu_name: "NVIDIA L4".to_string(), - }; - let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); - assert!(batch.unwrap_or(0.0) > 230.0, "L4 should support DQN batch > 230, got {:?}", batch); - } - - #[test] - fn test_dynamic_batch_size_h100() { - // H100 has 80GB VRAM — should hit the 8192 ceiling - let budget = crate::hyperopt::HardwareBudget { - gpu_memory_mb: 81_920, - gpu_name: "NVIDIA H100".to_string(), - }; - let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); - assert!((batch.unwrap_or(0.0) - 8192.0).abs() < 1.0, "H100 should hit 8192 ceiling, got {:?}", batch); - } - - // ── C2 Overhaul Smoke Tests ───────────────────────────────────────── - - /// Verify DQN action space is 5 exposure levels (not 45 factored actions). - #[test] - fn test_c2_dqn_default_num_actions_is_5() { - let config = crate::dqn::DQNConfig::default(); - assert_eq!(config.num_actions, 5, "DQN default must be 5 exposure-level actions"); - } - - /// Verify 5 exposure indices produce 5 distinct exposure levels. - #[test] - fn test_c2_five_actions_produce_distinct_exposures() { - use crate::dqn::action_space::ExposureLevel; - use crate::dqn::order_router::OrderRouter; - - let actions: Vec<_> = (0..5) - .filter_map(|idx| ExposureLevel::from_index(idx).ok()) - .map(|e| OrderRouter::route_default(e)) - .collect(); - - assert_eq!(actions.len(), 5); - - let unique: std::collections::HashSet<_> = actions.iter().map(|a| a.exposure).collect(); - assert_eq!(unique.len(), 5, "All 5 exposure levels must be distinct"); - } - - /// Verify hyperopt search space is 29D (C4: sharpe_weight, L2: branch_hidden_dim). - #[test] - fn test_c3_search_space_is_27d() { - let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 30, "Search space must be 30D (C6: gradient_accumulation_steps added)"); - - let names = crate::hyperopt::adapters::dqn::DQNParams::param_names(); - assert_eq!(names.len(), 30); - assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)"); - assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)"); - assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)"); - assert!(!names.contains(&"curiosity_weight"), "curiosity_weight must not be in search space"); - assert!(!names.contains(&"noisy_epsilon_floor"), "noisy_epsilon_floor must not be in search space"); - } - - /// Verify noisy_epsilon_floor is fixed to 0.10 (prevents action collapse). - #[test] - fn test_noisy_epsilon_floor_fixed() { - let params = crate::hyperopt::adapters::dqn::DQNParams::default(); - assert!( - (params.noisy_epsilon_floor - 0.10).abs() < 1e-6, - "noisy_epsilon_floor must default to 0.10 (prevents action collapse)" - ); - } - - /// Verify exploration params are fixed after C2 cleanup. - #[test] - fn test_c2_exploration_params_fixed() { - let params = crate::hyperopt::adapters::dqn::DQNParams::default(); - assert!( - params.curiosity_weight.abs() < f64::EPSILON, - "curiosity_weight must be fixed at 0.0" - ); - assert!( - (params.noisy_epsilon_floor - 0.10).abs() < 1e-6, - "noisy_epsilon_floor must be fixed at 0.10" - ); - assert!( - params.count_bonus_coefficient.abs() < f64::EPSILON, - "count_bonus_coefficient must be fixed at 0.0" - ); - - // Roundtrip through from_continuous should preserve fixed values - let continuous = params.to_continuous(); - let recovered = crate::hyperopt::adapters::dqn::DQNParams::from_continuous(&continuous).unwrap(); - assert!( - recovered.curiosity_weight.abs() < f64::EPSILON, - "curiosity_weight must remain 0.0 after roundtrip" - ); - assert!( - recovered.count_bonus_coefficient.abs() < f64::EPSILON, - "count_bonus_coefficient must remain 0.0 after roundtrip" - ); - } -} diff --git a/crates/ml/src/trainers/dqn/trainer/action.rs b/crates/ml/src/trainers/dqn/trainer/action.rs new file mode 100644 index 000000000..2f381196a --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/action.rs @@ -0,0 +1,396 @@ +//! DQN Trainer — Action selection, routing, and fill simulation + +use anyhow::Result; +use candle_core::Tensor; +use tracing::{debug, info}; + +use super::DQNTrainer; +use crate::dqn::action_space::{ExposureLevel, FactoredAction}; +use crate::dqn::TradingState; +use crate::dqn::mixed_precision::training_dtype; +use crate::dqn::order_router::OrderRouter; +use ml_core::fill_simulator::FillResult; + +impl DQNTrainer { + /// Select action using epsilon-greedy + pub(crate) async fn select_action(&self, state: &TradingState) -> Result { + let _agent = self.agent.read().await; + + // Convert state to tensor with tensor core alignment padding + let state_vec = state.to_vector(); + let raw_dim = state_vec.len(); + let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); + let padded: Vec = if aligned_dim > raw_dim { + let mut v = state_vec.to_vec(); + v.resize(aligned_dim, 0.0); + v + } else { + state_vec.to_vec() + }; + let state_tensor = Tensor::new(&*padded, &self.device) + .map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))? + .unsqueeze(0)?; // Add batch dimension + + // Get Q-values (epsilon-greedy handled by agent internally) + let action_idx = self.epsilon_greedy_action(&state_tensor).await?; + + let exposure = ExposureLevel::from_index(action_idx) + .map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?; + // Phase C: Use smart routing with trainer's spread/vol EMAs + Ok(self.route_action(exposure, self.hyperparams.avg_spread as f32)) + } + + /// Phase C: Route an exposure-level action using smart order routing. + /// + /// Uses the trainer's running spread/volatility EMAs to determine + /// optimal order type (Market/Limit/IoC) and urgency (Patient/Normal/Aggressive). + /// The DQN selects exposure; OrderRouter selects execution strategy. + pub(crate) fn route_action(&self, exposure: ExposureLevel, current_spread: f32) -> FactoredAction { + OrderRouter::route( + exposure, + current_spread, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ) + } + + /// Phase C: Simulate order fill and return result. + /// + /// Returns (action, fill_result) — if not filled, action is overridden to Flat + /// so the agent learns that limit orders in certain conditions don't execute. + pub(crate) fn simulate_fill( + &self, + action: FactoredAction, + step: usize, + ) -> (FactoredAction, FillResult) { + let normalized_vol = if self.median_vol > 0.0 { + (self.vol_ema / self.median_vol) as f32 + } else { + 1.0 + }; + let spread_bps = self.hyperparams.avg_spread * 10000.0; // fractional → bps + + let fill_result = self.fill_simulator.simulate_fill( + action.order, + action.urgency, + normalized_vol, + spread_bps, + step, + action.exposure as usize, + ); + + if fill_result.filled { + (action, fill_result) + } else { + // Order didn't fill — position stays unchanged (Flat action, no trade) + (OrderRouter::route_default(ExposureLevel::Flat), fill_result) + } + } + + /// Select actions for a batch of states (GPU-optimized) + /// + /// This method reduces GPU kernel launches by batching all action selections + /// into a single forward pass. Provides 125× reduction in kernel launches + /// compared to sequential select_action() calls. + /// + /// # Performance Impact + /// - Single GPU kernel launch for entire batch (vs. one per sample) + /// - Reduced CPU-GPU synchronization overhead + /// - Better GPU utilization through larger batch sizes + /// + /// # Arguments + /// * `states` - Slice of TradingState objects to process + /// + /// # Returns + /// Vector of TradingAction decisions (same order as input states) + pub(crate) async fn select_actions_batch(&mut self, states: &[TradingState]) -> Result> { + if states.is_empty() { + return Ok(Vec::new()); + } + + let agent = self.agent.read().await; + let batch_size = states.len(); + + // Get state dimension from first state, then align for tensor cores + let first_vec = states + .first() + .map(|s| s.to_vector()) + .ok_or_else(|| anyhow::anyhow!("Empty states slice"))?; + let raw_state_dim = first_vec.len(); + let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &self.device); + let pad = aligned_dim - raw_state_dim; + + // Pre-allocate flat buffer, zero-padding each state to aligned dimension + let mut flat_states = Vec::with_capacity(batch_size * aligned_dim); + flat_states.extend_from_slice(&first_vec); + flat_states.extend(std::iter::repeat_n(0.0_f32, pad)); + + for (i, state) in states.iter().enumerate().skip(1) { + let vec = state.to_vector(); + if vec.len() != raw_state_dim { + return Err(anyhow::anyhow!( + "State {} dimension mismatch: expected {}, got {}", + i, + raw_state_dim, + vec.len() + )); + } + flat_states.extend_from_slice(&vec); + flat_states.extend(std::iter::repeat_n(0.0_f32, pad)); + } + + // Create batched tensor directly from flat buffer + let batch_tensor = Tensor::from_vec(flat_states, (batch_size, aligned_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Failed to cast batched state tensor to training dtype: {}", e))?; + + // FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor. + // Previously used get_epsilon() which returns the decayed epsilon (0.0 with noisy nets), + // making the batch path have ZERO random exploration — root cause of action collapse. + let base_epsilon = agent.get_effective_epsilon() as f64; + let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon); + let epsilon = adjusted_epsilon as f32; + + debug!("Epsilon: base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon); + + // Single forward pass for all samples (GPU-optimized). + // For RegimeConditional, forward() now blends all 3 heads via regime masks. + let batch_q_values = agent + .forward(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?; + + // Branching DQN: get per-branch Q-values while agent lock is held. + // Returns (exposure [batch,5], order [batch,3], urgency [batch,3]). + let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = + if self.hyperparams.use_branching { + agent + .batch_branching_q_values(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))? + } else { + None + }; + + drop(agent); // Release lock early + + // Fused GPU epsilon-greedy: argmax + RNG in a single CUDA kernel launch. + // Eliminates the intermediate GPU→CPU sync from Candle's argmax(). + // Lazy-init the GPU action selector on first call + { + if self.gpu_action_selector.is_none() && self.device.is_cuda() { + let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( + &self.device, + self.hyperparams.batch_size.max(batch_size).max(8192), + 0xDEAD_BEEF_CAFE_u64, + ).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?; + info!("GPU action selector initialized for select_actions_batch"); + self.gpu_action_selector = Some(selector); + } + + let selector = self.gpu_action_selector.as_mut() + .ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?; + + // Branching path: per-branch epsilon-greedy → factored indices (0-44) + let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { + selector + .select_actions_branching(q_exp, q_ord, q_urg, epsilon) + .map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))? + } else { + // Non-branching: exposure-only epsilon-greedy → GPU route to factored + let exposure_tensor = selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; + + selector.route_exposure_to_factored( + &exposure_tensor, + batch_size, + self.hyperparams.avg_spread as f32, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? + }; + + // Per-element scalar readback (no to_vec1) + // narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar + let mut actions = Vec::with_capacity(batch_size); + for i in 0..batch_size { + let idx = factored_tensor.narrow(0, i, 1) + .and_then(|t| t.squeeze(0)) + .and_then(|t| t.to_scalar::()) + .map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?; + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?; + actions.push(action); + } + Ok(actions) + } + } + + /// GPU-optimized batch action selection using pre-built state tensor. + /// + /// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch). + /// GPU-batched action selection with optional fused routing + fill simulation. + /// + /// Returns `(actions, gpu_handled_fill)`: + /// - `gpu_handled_fill = true`: actions are post-fill (routed + fill-checked by GPU kernel). + /// Caller must NOT apply CPU `route_action()` or `simulate_fill()` — already done. + /// - `gpu_handled_fill = false`: actions have basic routing only. Caller should apply + /// CPU routing + fill as before. + /// + /// The fused kernel (`epsilon_greedy_routed`) is used when: + /// 1. GPU action selector is available (CUDA device) + /// 2. Not using branching DQN (branching learns order type via network heads) + /// 3. Median volatility > 0 (fill simulation requires vol context) + pub(crate) async fn select_actions_batch_gpu( + &mut self, + batch_tensor: &Tensor, + batch_start: usize, + ) -> Result<(Vec, bool)> { + let batch_size = batch_tensor.dims()[0]; + if batch_size == 0 { + return Ok((Vec::new(), false)); + } + + let agent = self.agent.read().await; + + // FIX: Use get_effective_epsilon() which respects noisy_epsilon_floor. + // Same fix as select_actions_batch() — previously used get_epsilon() which + // returned 0.0 with noisy nets, causing zero exploration in GPU batch path. + let base_epsilon = agent.get_effective_epsilon() as f64; + let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon); + let epsilon = adjusted_epsilon as f32; + debug!("Epsilon (GPU path): base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon); + + // Single forward pass — tensor already on GPU, no construction needed + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("GPU batched forward pass failed: {}", e))?; + + // C2 FIX: Count bonus removed from Q-value computation. + // Noisy nets are the sole exploration mechanism during training. + // Count bonus kept for diversity metrics only (record_action tracking). + + // Branching DQN: get per-branch Q-values if branching mode is active. + // Uses unified dispatch that supports both Standard and RegimeConditional. + // For RC, regime classification masks blend per-branch Q-values from all 3 heads. + let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = if self.hyperparams.use_branching { + agent + .batch_branching_q_values(batch_tensor) + .map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))? + } else { + None + }; + + drop(agent); + + // Fused GPU epsilon-greedy: argmax + RNG + routing in CUDA kernels. + // Lazy-init GPU action selector on first call. + if self.gpu_action_selector.is_none() && self.device.is_cuda() { + let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( + &self.device, + self.hyperparams.batch_size.max(batch_size).max(8192), + 0xDEAD_BEEF_CAFE_u64, + ).map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?; + info!("GPU action selector initialized for select_actions_batch_gpu"); + self.gpu_action_selector = Some(selector); + } + + let selector = self.gpu_action_selector.as_mut() + .ok_or_else(|| anyhow::anyhow!("GPU action selector not initialized on non-CUDA device"))?; + + // Use fused routing+fill kernel when conditions allow: + // - Not branching (branching learns order type via network heads) + // - Median vol > 0 (fill simulation needs vol context) + let use_routed = !self.hyperparams.use_branching && self.median_vol > 0.0; + + // Select actions via fused kernel — one launch, no intermediate GPU→CPU sync + let factored_tensor = if self.hyperparams.use_branching { + if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { + // Branching: 3-head epsilon-greedy → factored (0-44) directly + selector + .select_actions_branching(q_exp, q_ord, q_urg, epsilon) + .map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))? + } else { + // Branching Q forward failed — exposure-only → route on GPU + let exposure_tensor = selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; + selector.route_exposure_to_factored( + &exposure_tensor, batch_size, + self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, + self.vol_ema as f32, self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? + } + } else if use_routed { + // Fused: epsilon-greedy + routing + fill simulation in one kernel + let spread = self.hyperparams.avg_spread as f32; + let spread_bps = (self.hyperparams.avg_spread * 10000.0) as f32; + selector + .select_actions_routed( + &batch_q_values, epsilon, batch_size, 5, + batch_start as i32, + spread, spread, + self.vol_ema as f32, self.median_vol as f32, + spread_bps, 0.85, 0.30, 0.80, 0.50, 0.50, + ) + .map_err(|e| anyhow::anyhow!("GPU routed action selection failed: {e}"))? + } else { + // Exposure-only → route on GPU + let exposure_tensor = selector + .select_actions(&batch_q_values, epsilon, batch_size, 5) + .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?; + selector.route_exposure_to_factored( + &exposure_tensor, batch_size, + self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, + self.vol_ema as f32, self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? + }; + + // Per-element scalar readback (no to_vec1) + // narrow(0,i,1) → shape [1]; squeeze(0) → scalar [] for to_scalar + let mut actions = Vec::with_capacity(batch_size); + for i in 0..batch_size { + let idx = factored_tensor.narrow(0, i, 1) + .and_then(|t| t.squeeze(0)) + .and_then(|t| t.to_scalar::()) + .map_err(|e| anyhow::anyhow!("Factored index readback [{i}]: {e}"))?; + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {idx}: {e}"))?; + actions.push(action); + } + Ok((actions, use_routed)) + } + + /// Epsilon-greedy action selection for single-step inference. + /// + /// This is batch_size=1 — the overhead of a CUDA kernel launch (~5us) exceeds + /// the benefit of fusing argmax+RNG for a single element. Candle's argmax + + /// to_scalar is already minimal for this path. Keep on CPU. + pub(crate) async fn epsilon_greedy_action(&self, state: &Tensor) -> Result { + use rand::Rng; + + let epsilon = self.get_epsilon().await? as f32; + let mut rng = rand::thread_rng(); + + if rng.gen::() < epsilon { + // Random action (exploration) over 5 exposure levels + Ok(rng.gen_range(0..5)) + } else { + // Greedy action (exploitation) - use actual Q-network + let agent = self.agent.read().await; + let q_values = agent.forward(state)?; + + // GPU-native argmax — single u32 scalar transfer instead of full Q-value vector + let best_action = q_values + .argmax(1) + .and_then(|t| t.squeeze(0)) + .and_then(|t| t.to_scalar::()) .map(|v| v as usize) + .ok() + .unwrap_or(2); // Default to HOLD (index 2) on error + + Ok(best_action) + } + } + +} diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs new file mode 100644 index 000000000..25ee4106d --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -0,0 +1,732 @@ +//! DQN Trainer constructor — `new_internal` and init helpers. + +use std::collections::VecDeque; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use candle_core::Device; +use ml_core::fill_simulator::FillSimulator; +use risk::drawdown_monitor::DrawdownMonitor; +use risk::safety::position_limiter::HybridPositionLimiter; +use risk::safety::PositionLimiterConfig; +use rust_decimal::Decimal; +use tokio::sync::RwLock; +use tracing::info; + +use crate::dqn::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; +use crate::dqn::curiosity::CuriosityModule; +use crate::dqn::dqn::{DQN, DQNConfig}; +use crate::dqn::logging::{LoggingConfig, MetricsAggregator}; +use crate::dqn::portfolio_tracker::PortfolioTracker; +use crate::dqn::regime_conditional::RegimeConditionalDQN; +use crate::dqn::reward::{RewardConfig, RewardFunction}; +use crate::features::microstructure_features::*; +use crate::labeling::triple_barrier::TripleBarrierEngine; +use crate::memory_optimization::auto_batch_size::{AutoBatchSizer, BatchSizeConfig}; +use crate::trainers::TargetUpdateMode; +use crate::TrainingMetrics; +use super::super::config::{DQNAgentType, DQNHyperparameters}; +use super::DQNTrainer; + +impl DQNTrainer { + pub(crate) fn new_internal(mut hyperparams: DQNHyperparameters, debug_logging: bool, override_device: Option) -> Result { + // Validate batch size is non-zero + if hyperparams.batch_size == 0 { + return Err(anyhow::anyhow!( + "Batch size must be greater than 0, got: {}", + hyperparams.batch_size + )); + } + + // WAVE 26 P2.2: Validate gradient_accumulation_steps > 0 + if hyperparams.gradient_accumulation_steps == 0 { + return Err(anyhow::anyhow!( + "gradient_accumulation_steps must be greater than 0, got: {}", + hyperparams.gradient_accumulation_steps + )); + } + + // Pre-compute hidden dims to get accurate model size for batch sizing. + // Align input_dim to 8 so the log matches the actual model dimensions. + // (device not yet created, so use the formula directly — CUDA always aligns) + let ofi_pre = hyperparams.mbp10_data_dir.is_some(); + let input_dim: usize = if ofi_pre { 56 } else { 48 }; // (53+7)&!7=56, (45+7)&!7=48 + let output_dim: usize = 5; + let hidden_dims: Vec = match hyperparams.hidden_dim_base { + Some(base) => { + let b = crate::cuda_pipeline::align_to_tensor_cores(base); + vec![b, b] // Constant-width: no tapering, no silently-discarded narrow layer + } + None => { + let caps = crate::gpu::capabilities::cached_capabilities(); + let base = crate::gpu::memory_profile::resolve_hidden_dim_base( + caps.free_vram_mb, + ); + let b = crate::cuda_pipeline::align_to_tensor_cores(base); + vec![b, b] // Constant-width: no tapering + } + }; + + // Compute accurate model size from actual network dimensions + let full_dims: Vec = std::iter::once(input_dim) + .chain(hidden_dims.iter().copied()) + .chain(std::iter::once(output_dim)) + .collect(); + let param_count = crate::gpu::memory_profile::network_param_count(&full_dims); + // FP32 params + AdamW state (2x for momentum/variance) + target network copy = ~4x + let model_overhead_mb = (param_count as f64 * 4.0 * 4.0) / (1024.0 * 1024.0); + info!( + "DQN network: {:?} → {} params, {:.1} MB overhead", + full_dims, param_count, model_overhead_mb + ); + + // Dynamic batch sizing: scale UP for larger GPUs, cap DOWN for smaller ones. + // Uses HardwareBudget for consistent sizing across DQN/PPO. + const STATIC_MAX_BATCH_SIZE: usize = 8192; + let max_safe_batch = match AutoBatchSizer::new() { + Ok(sizer) => { + let config = BatchSizeConfig { + model_memory_mb: model_overhead_mb, + safety_margin: 0.15, + ..BatchSizeConfig::default() + }; + let safe = sizer.max_safe_batch_size(&config); + info!( + "AutoBatchSizer: GPU VRAM ceiling = {} (configured: {})", + safe, hyperparams.batch_size + ); + safe + } + Err(e) => { + info!( + "AutoBatchSizer unavailable ({}), using static cap: {}", + e, STATIC_MAX_BATCH_SIZE + ); + STATIC_MAX_BATCH_SIZE + } + }; + + // Cap to VRAM ceiling from AutoBatchSizer (no separate scale-UP — + // AutoBatchSizer already accounts for model size and available VRAM) + if hyperparams.batch_size > max_safe_batch { + info!( + "DQN batch_size capped from {} → {} (VRAM ceiling)", + hyperparams.batch_size, max_safe_batch + ); + hyperparams.batch_size = max_safe_batch; + } + + // Use override device if provided (hyperopt shares one CUDA context), + // otherwise auto-detect GPU + let device = if let Some(dev) = override_device { + dev + } else { + Device::cuda_if_available(0) + .map_err(|e| anyhow::anyhow!("Failed to initialize device: {}", e))? + }; + + // Dynamic replay buffer sizing: scale replay capacity to available VRAM. + // Only activates when replay_buffer_vram_fraction > 0 and GPU is detected. + // Also computes the PER memory budget from actual VRAM — no hardcoded caps. + let original_buffer_size = hyperparams.buffer_size; // Save before AutoReplaySizer mutates it + let mut per_max_memory_bytes: usize = 4 * 1024 * 1024 * 1024; // CPU fallback: 4 GB + if hyperparams.replay_buffer_vram_fraction > 0.0 && device.is_cuda() { + use ml_core::memory_optimization::detect_gpu_hardware; + match detect_gpu_hardware() { + Ok(hw) => { + let raw_sd = if hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; + let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &device); + let replay_cfg = hw.optimal_replay_config( + aligned_sd, + hyperparams.replay_buffer_vram_fraction, + ); + per_max_memory_bytes = replay_cfg.per_max_buffer_bytes; + if replay_cfg.capacity != hyperparams.buffer_size { + info!( + "AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%, PER budget={:.0}MB)", + hyperparams.buffer_size, + replay_cfg.capacity, + hw.free_memory_mb, + hyperparams.replay_buffer_vram_fraction * 100.0, + per_max_memory_bytes as f64 / (1024.0 * 1024.0), + ); + hyperparams.buffer_size = replay_cfg.capacity; + } + } + Err(e) => { + info!( + "AutoReplaySizer unavailable ({}), using static buffer_size: {}", + e, hyperparams.buffer_size + ); + } + } + } else if device.is_cuda() { + // No auto-sizer, but still compute PER budget from VRAM + use ml_core::memory_optimization::detect_gpu_hardware; + if let Ok(hw) = detect_gpu_hardware() { + per_max_memory_bytes = hw.per_max_buffer_bytes(); + } else { + // GPU detection failed — keep CPU default (4 GB) + } + } else { + // CPU device — keep default 4 GB PER budget + } + + info!( + "Initializing DQN trainer on device: {:?}, using 5 exposure actions + OrderRouter", + if device.is_cuda() { "CUDA GPU" } else { "CPU" }, + ); + + // Auto-detect mixed precision capability based on GPU architecture + let mixed_precision_detected = if device.is_cuda() { + match crate::memory_optimization::auto_batch_size::detect_gpu_memory() { + Ok((_total, _free, ref name)) => { + let detected = crate::dqn::mixed_precision::detect_from_gpu_name(name); + match &detected { + Some(c) => info!("GPU mixed precision: {:?} enabled (GPU: {})", c.dtype, name), + None => info!("GPU mixed precision: disabled (GPU: {})", name), + } + detected + } + Err(_) => None, + } + } else { + None + }; + + // Create DQN configuration + // 42-feature architecture: OHLCV, technical, patterns, volume, time, statistical, regime + // Portfolio features (3) are populated via PortfolioTracker → 45 total state_dim + // With MBP-10 OFI features: +8 OFI features → 53 total + // + // Tensor core alignment: state_dim is rounded up to the next multiple of 8 + // (53→56, 45→48) so that cuBLAS dispatches BF16 HMMA instructions instead + // of falling back to scalar FMA. The extra columns are zero-padded at the + // data pipeline boundaries (GpuPreloadedData and train_batch CPU path). + let ofi_enabled = hyperparams.mbp10_data_dir.is_some(); + let raw_state_dim = if ofi_enabled { 53 } else { 45 }; + let state_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &device); + let config = DQNConfig { + state_dim, + num_actions: 5, // 5 exposure levels (Short100, Short50, Flat, Long50, Long100) + hidden_dims, + learning_rate: hyperparams.learning_rate, + gamma: hyperparams.gamma as f32, + epsilon_start: hyperparams.epsilon_start as f32, + epsilon_end: hyperparams.epsilon_end as f32, + epsilon_decay: hyperparams.epsilon_decay as f32, + replay_buffer_capacity: hyperparams.buffer_size, + collapse_warmup_capacity: original_buffer_size, + batch_size: hyperparams.batch_size, + min_replay_size: hyperparams.min_replay_size.min(hyperparams.buffer_size), // Cap at buffer_size to prevent deadlock + target_update_freq: hyperparams.target_update_frequency, // Use hyperparameter instead of hardcoded 1000 + use_double_dqn: true, + use_huber_loss: hyperparams.use_huber_loss, + huber_delta: hyperparams.huber_delta as f32, + leaky_relu_alpha: 0.01, // Standard LeakyReLU alpha (prevents dead neurons) + gradient_clip_norm: hyperparams.gradient_clip_norm.unwrap_or(10.0), // Wave 11 Bug #1 fix: Dynamic clipping + + // WAVE 16 (Agent 36): Target update configuration + tau: hyperparams.tau, + tau_final: hyperparams.tau * 0.1, // Anneal to 10% of base tau + tau_anneal_steps: 100_000, + use_soft_updates: matches!(hyperparams.target_update_mode, TargetUpdateMode::Soft), + + // Rainbow DQN warmup period + warmup_steps: hyperparams.warmup_steps, + + // PER configuration + initial_capital: hyperparams.initial_capital as f64, + use_per: hyperparams.use_per, + use_gpu_replay_buffer: hyperparams.use_gpu_replay_buffer, + per_alpha: hyperparams.per_alpha, + per_beta_start: hyperparams.per_beta_start, + per_beta_max: 1.0, + per_beta_annealing_steps: hyperparams.epochs * 2000, // ~2000 steps/epoch (130k bars / ~64 batch_size) + per_max_memory_bytes, + + // Wave 2.1: Dueling Networks (ENABLED BY DEFAULT - Wave 6.4) + use_dueling: hyperparams.use_dueling, + dueling_hidden_dim: hyperparams.dueling_hidden_dim, + + // Wave 2.2: Multi-Step Returns (N-step TD) (ENABLED BY DEFAULT - Wave 6.4) + n_steps: hyperparams.n_steps, // Default: 3 (Rainbow DQN standard) + + // Wave 2.3: Distributional RL (C51) (ENABLED BY DEFAULT - Wave 6.4) + use_distributional: hyperparams.use_distributional, // Default: enabled (C51 distributional RL) + num_atoms: hyperparams.num_atoms, // Rainbow DQN standard: 51 atoms + v_min: hyperparams.v_min as f32, // Minimum value for distribution support + v_max: hyperparams.v_max as f32, // Maximum value for distribution support + + // Wave 2.4: Noisy Networks for Exploration (ENABLED BY DEFAULT - Wave 6.4) + use_noisy_nets: hyperparams.use_noisy_nets, // Default: enabled (replaces epsilon-greedy) + noisy_sigma_init: hyperparams.noisy_sigma_init, // Rainbow DQN standard: 0.5 + + // BUG #37 FIX: Q-value clipping (prevents step-level explosions) + enable_q_value_clipping: true, + q_value_clip_min: -500.0, + q_value_clip_max: 500.0, + + // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold (from hyperparams) + gradient_collapse_multiplier: hyperparams.gradient_collapse_multiplier, + gradient_collapse_patience: hyperparams.gradient_collapse_patience, + + use_cql: hyperparams.use_cql, + cql_alpha: hyperparams.cql_alpha, + use_iqn: hyperparams.use_qr_dqn, // Controlled by hyperopt + iqn_num_quantiles: hyperparams.num_quantiles, // Controlled by hyperopt + iqn_kappa: hyperparams.qr_kappa as f32, // Controlled by hyperopt (f64→f32) + iqn_embedding_dim: 64, // Fixed (not in search space) + use_branching: hyperparams.use_branching, + branch_hidden_dim: hyperparams.branch_hidden_dim, + use_regime_conditioning: true, // Always enable per-regime IS weights for branching loss + use_cvar_action_selection: false, + cvar_alpha: 0.05, + + #[allow(clippy::cast_possible_truncation)] + minimum_profit_factor: hyperparams.minimum_profit_factor as f32, + weight_decay: hyperparams.weight_decay, + dropout_rate: if hyperparams.enable_dropout_scheduler { hyperparams.dropout_initial } else { 0.0 }, + mixed_precision: hyperparams.mixed_precision.clone().or(mixed_precision_detected), + entropy_coefficient: hyperparams.entropy_coefficient.unwrap_or(0.01), + noisy_epsilon_floor: hyperparams.noisy_epsilon_floor.unwrap_or(0.0) as f32, // C2: NoisyNet handles exploration + use_count_bonus: hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, // C3 FIX: enable when coefficient > 0 + count_bonus_coefficient: hyperparams.count_bonus_coefficient.unwrap_or(0.0), + ..DQNConfig::default() + }; + + // Extract curiosity dims before config is moved into the agent + let curiosity_market_dim = config.curiosity_market_dim; + let curiosity_hidden_dim = config.curiosity_hidden_dim; + + // Create DQN agent + let agent = if hyperparams.enable_regime_qnetwork { + info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)"); + info!(" - Regime detection: ADX (raw index 40) + CUSUM direction (raw index 41)"); + info!(" - Classification: Trending (ADX>0.25), Volatile (ADX≤0.25 & |CUSUM|>0.7), Ranging (otherwise)"); + let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone()) + .map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?; + DQNAgentType::RegimeConditional(regime_agent) + } else { + info!("Creating standard DQN with single Q-network head"); + let standard_agent = DQN::new_on_device(config, device.clone()) + .map_err(|e| anyhow::anyhow!("Failed to create DQN agent: {}", e))?; + DQNAgentType::Standard(standard_agent) + }; + + + // Initialize portfolio tracker with $100k starting capital and 1 basis point spread + // Bug #2 fix: Portfolio features were hardcoded as [0.0, 0.0, 0.0] at line 1528 + let portfolio_tracker = PortfolioTracker::new( + hyperparams.initial_capital, // P2-A: Configurable capital + 0.0001, // 1 basis point spread (0.01%) + hyperparams.cash_reserve_percent, // Cash reserve requirement + ); + + // Initialize reward function with hyperparameter-driven configuration + // WAVE 10-A9 FIX: Wire hold_penalty_weight from hyperparameters to RewardConfig + // BUG #17 FIX: Add normalization and percentage-based P&L (enabled by default) + let reward_config = RewardConfig { + pnl_weight: Decimal::ONE, + risk_weight: Decimal::try_from(0.1).unwrap_or(Decimal::ZERO), + cost_weight: Decimal::ONE, // Bug #2 fix: 100% transaction cost weight (was 0.05, 20x too low) + hold_reward: Decimal::ZERO, // Flat position = no edge = zero reward (was +0.001, 20x trade PnL) + movement_threshold: Decimal::try_from(hyperparams.movement_threshold) + .unwrap_or(Decimal::ZERO), + hold_penalty_weight: Decimal::try_from(hyperparams.hold_penalty_weight) + .unwrap_or(Decimal::ZERO), // CRITICAL FIX + diversity_weight: Decimal::try_from(-0.1).unwrap_or(Decimal::ZERO), + enable_normalization: !hyperparams.use_dsr, // DSR replaces EMA normalizer + use_percentage_pnl: true, // Bug #17: Use percentage returns for scale-invariance + circuit_breaker_config: CircuitBreakerConfig::default(), + triple_barrier_profit_bonus: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO), + triple_barrier_stop_penalty: Decimal::try_from(0.5).unwrap_or(Decimal::ZERO), + sharpe_weight: Decimal::ZERO, // WAVE 26 P1.3: Disabled by default + sharpe_window: 20, // WAVE 26 P1.3: Standard 20-period window + use_dsr: hyperparams.use_dsr, + dsr_eta: hyperparams.dsr_eta, + initial_capital: hyperparams.initial_capital as f64, + }; + let reward_fn = RewardFunction::new_with_debug(reward_config, debug_logging)?; + + // WAVE 1.1: Initialize triple barrier engine (max 1000 active trackers) + let triple_barrier = Arc::new(RwLock::new(TripleBarrierEngine::new(1000))); + info!("Triple barrier engine initialized with 1000 max trackers"); + + // WAVE 16S: Initialize Kelly optimizer if enabled + let kelly_optimizer = if hyperparams.enable_kelly_sizing { + use crate::risk::kelly_optimizer::{KellyCriterionOptimizer, KellyOptimizerConfig}; + let kelly_config = KellyOptimizerConfig { + max_fraction: hyperparams.kelly_max_fraction, + min_fraction: 0.01, + lookback_period: 252, + confidence_threshold: 0.6, + volatility_adjustment: true, + drawdown_protection: true, + }; + let optimizer = KellyCriterionOptimizer::new(kelly_config) + .map_err(|e| anyhow::anyhow!("Failed to create Kelly optimizer: {}", e))?; + info!("Kelly optimizer enabled (fractional={}, max={})", + hyperparams.kelly_fractional, hyperparams.kelly_max_fraction); + Some(Arc::new(optimizer)) + } else { + None + }; + + // Wave 16 Portfolio Features: Initialize action masking, entropy regularization, and stress testing + let enable_action_masking = hyperparams.enable_action_masking; + let max_position = hyperparams.max_position_absolute; // BLOCKER #2: Use hyperopt-tunable position limit + + // Entropy regularization: SAC-style computed directly on Q-values in DQN::compute_loss_internal + if hyperparams.enable_entropy_regularization { + let coeff = hyperparams.entropy_coefficient.unwrap_or(0.01); + info!("Entropy regularization enabled (coefficient={coeff:.4}, applied to Q-value softmax in loss)"); + } + + // Multi-asset portfolio tracking (disabled -- single-asset is the current production mode) + // + // When expanding to multi-asset trading: + // 1. Add `enable_multi_asset: bool` to DQNHyperparams (default false). + // 2. Initialize MultiAssetPortfolioTracker here when the flag is set. + // 3. Wire portfolio state into the DQN observation: expand state_dim to + // include per-asset position, PnL, and correlation features so the + // agent can learn cross-asset hedging and allocation. + let multi_asset_portfolio: Option> = None; + + // Stress testing for robustness validation + // Initialized as None here; call init_stress_tester() after construction + // to resolve the circular dependency (DQNStressTester needs a DQNTrainer). + let stress_tester: Option> = None; + + if enable_action_masking { + info!( + "Action masking enabled (max_position=±{:.1}, 30-50% filtering expected)", + max_position + ); + } else { + info!("Action masking disabled (all 5 exposure levels available)"); + } + + // Wave 16 Core Risk Features: Initialize drawdown monitor, position limiter, circuit breaker + // These are ALWAYS enabled by default for production safety + + // 1. Drawdown Monitor (15% max drawdown, alerts at 10%, 12.5%, 15%) + let drawdown_monitor = { + // DrawdownMonitor will be configured in first training step + // Config will be applied via async configure_alerts() in train_epoch + info!("Drawdown monitor enabled (thresholds: 10%, 12.5%, 15%)"); + Some(Arc::new(DrawdownMonitor::new())) + }; + + // 2. Position Limiter (3-tier limits: ±10.0 absolute, 1M notional, 10% concentration) + let position_limiter = { + let config = PositionLimiterConfig { + enabled: true, + cache_ttl: Duration::from_secs(60), + rpc_check_threshold_percent: 0.8, + max_position_per_symbol: 10.0, // ±10.0 absolute position limit + max_order_value: 1_000_000.0, // $1M notional limit + max_daily_loss: 0.10, // 10% concentration limit + }; + let limiter = HybridPositionLimiter::new(config); + info!("Position limiter enabled (abs=±10.0, notional=$1M, concentration=10%)"); + Some(Arc::new(limiter)) + }; + + // 3. Circuit Breaker (5 consecutive failures, 60s cooldown) + let circuit_breaker = { + let config = CircuitBreakerConfig { + failure_threshold: 5, + success_threshold: 3, + timeout_duration: Duration::from_secs(60), + half_open_max_calls: 2, + }; + let breaker = CircuitBreaker::new(config); + info!("Circuit breaker enabled (threshold=5 failures, cooldown=60s)"); + Some(Arc::new(breaker)) + }; + + // WAVE 24: Capture patience before hyperparams is moved + let early_stopping_patience = hyperparams.gradient_collapse_patience; + + // WAVE 26 P1: Initialize advanced DQN features + // P1.3: Sharpe ratio reward component + let sharpe_weight = hyperparams.sharpe_weight; + let sharpe_window = hyperparams.sharpe_window; + + // P1.6: Adaptive dropout scheduler + let dropout_scheduler = hyperparams.enable_dropout_scheduler.then(|| { + use crate::dqn::network::DropoutScheduler; + info!("Dropout scheduler enabled (initial={}, final={}, steps={})", + hyperparams.dropout_initial, hyperparams.dropout_final, hyperparams.dropout_anneal_steps); + DropoutScheduler::new( + hyperparams.dropout_initial, + hyperparams.dropout_final, + hyperparams.dropout_anneal_steps, + ) + }); + + // P1.7: Hindsight Experience Replay (HER) + let her_buffer = (hyperparams.her_ratio > 0.0) + .then(|| { + use crate::dqn::hindsight_replay::{HindsightReplayBuffer, HindsightReplayConfig, HindsightStrategy}; + use crate::dqn::prioritized_replay::PrioritizedReplayConfig; + let her_strategy = match hyperparams.her_strategy.as_str() { + "final" => HindsightStrategy::Final, + _ => HindsightStrategy::Future, // Default to Future + }; + let config = HindsightReplayConfig { + base_config: PrioritizedReplayConfig { + capacity: hyperparams.buffer_size, + ..Default::default() + }, + her_ratio: hyperparams.her_ratio, + her_strategy, + goal_dim: 1, // Single goal dimension for trading (target return) + k_future: 4, // Sample 4 future goals for Future strategy + batch_size: hyperparams.batch_size, + }; + info!("HER buffer enabled (ratio={}, strategy={:?}, capacity={})", + hyperparams.her_ratio, her_strategy, hyperparams.buffer_size); + HindsightReplayBuffer::new(config) + .map(Arc::new) + .map_err(|e| anyhow::anyhow!("Failed to create HER buffer: {}", e)) + }) + .transpose()?; + + // P1.9: Generalized Advantage Estimation (GAE) + let gae_calculator = hyperparams.enable_gae.then(|| { + use crate::dqn::gae::GAECalculator; + info!("GAE calculator enabled (lambda={}, gamma={})", + hyperparams.gae_lambda, hyperparams.gamma); + GAECalculator::new(hyperparams.gae_lambda, hyperparams.gamma) + }); + + // P1.11: Noisy network sigma scheduling + let noisy_sigma_scheduler = hyperparams.enable_noisy_sigma_scheduler.then(|| { + use crate::dqn::noisy_sigma_scheduler::NoisySigmaScheduler; + info!("Noisy sigma scheduler enabled (initial={}, final={}, steps={})", + hyperparams.noisy_sigma_initial, hyperparams.noisy_sigma_final, hyperparams.noisy_sigma_anneal_steps); + NoisySigmaScheduler::new( + hyperparams.noisy_sigma_initial, + hyperparams.noisy_sigma_final, + hyperparams.noisy_sigma_anneal_steps, + ) + }); + + // WAVE 26 P1.8: Initialize curiosity module if curiosity_weight > 0 + // Must be created BEFORE hyperparams and device are moved + let curiosity_module = (hyperparams.curiosity_weight > 0.0) + .then(|| { + CuriosityModule::new( + device.clone(), + 0.001, // Forward model learning rate + 2.0, // Max curiosity reward (clip to prevent noise exploitation) + curiosity_market_dim, + curiosity_hidden_dim, + 3, // Action categories (Short/Flat/Long) + ).map_err(|e| anyhow::anyhow!("Failed to create curiosity module: {}", e)) + }) + .transpose()?; + + // WAVE 26 P0.6: Initialize learning rate scheduler with warmup + // Must be created BEFORE hyperparams is moved + let lr_scheduler = { + use super::super::lr_scheduler::LRScheduler; + LRScheduler::new( + hyperparams.learning_rate, + hyperparams.warmup_steps, + hyperparams.lr_decay_type, + ) + }; + + // WAVE 44: Initialize n-step buffer if n_steps > 1 + let nstep_buffer = (hyperparams.n_steps > 1).then(|| { + info!("🎯 Multi-step returns ENABLED: n_steps={}, gamma={}", + hyperparams.n_steps, hyperparams.gamma); + crate::dqn::nstep_buffer::NStepBuffer::new( + hyperparams.n_steps, + hyperparams.gamma + ) + }); + + // Capture values before hyperparams is moved into Self + let initial_batch_size = hyperparams.batch_size; + let base_tau = hyperparams.tau; + + // GPU pipeline: pre-allocate staging buffers on CUDA devices + let buffer_pool = device.is_cuda().then(|| { + info!("GpuBufferPool: pre-allocated staging buffers (100k bars, 42 features, 4 targets)"); + crate::cuda_pipeline::GpuBufferPool::new(100_000, 42, 4) + }); + + // GPU pipeline: double-buffered loader for zero-downtime fold transitions + let double_buffer = device.is_cuda().then(|| { + crate::cuda_pipeline::double_buffer::DoubleBufferedLoader::new(device.clone()) + }); + + // Multi-GPU: auto-detect if multiple CUDA devices are available + let multi_gpu = crate::cuda_pipeline::multi_gpu::MultiGpuConfig::detect() + .unwrap_or(None); + if let Some(ref mg) = multi_gpu { + info!("Multi-GPU: {} devices detected, data parallelism enabled", mg.world_size); + } + + Ok(Self { + agent: Arc::new(RwLock::new(agent)), + hyperparams, + device, + metrics: Arc::new(RwLock::new(TrainingMetrics::new())), + loss_history: Vec::new(), + q_value_history: Vec::new(), + best_val_loss: f64::INFINITY, // Start with worst possible loss + val_data: Vec::new(), + val_loss_history: Vec::new(), + sharpe_history: Vec::new(), + best_sharpe: f64::NEG_INFINITY, // C4: Start with worst possible Sharpe + best_epoch: 0, + gradient_logging_step: 0, + collapse_warmup_buffer_size: original_buffer_size, + portfolio_tracker, + feature_stats: None, // WAVE 3 FIX #2: Start with None, collect stats in epochs 0-10 + recent_actions: VecDeque::with_capacity(100), + reward_fn, + + // WAVE 16S: Adaptive risk management + kelly_optimizer, + trade_history: VecDeque::with_capacity(500), + volatility_returns: VecDeque::with_capacity(20), // Use default instead of moved hyperparams + pnl_history: VecDeque::with_capacity(1000), + + // Wave 16 Portfolio Features + enable_action_masking, + max_position, + multi_asset_portfolio, + stress_tester, + + // Wave 16 Core Risk Features + drawdown_monitor, + position_limiter, + circuit_breaker, + + // WAVE 3.10: Microstructure feature calculators + micro_high_low_spread: HighLowSpread::default(), + micro_vw_spread: VolumeWeightedSpread::default(), + micro_tick_count: TickCount::default(), + micro_inter_arrival: InterArrivalTime::default(), + micro_buy_sell_imbalance: BuySellImbalance::default(), + micro_kyle_lambda: KyleLambda::default(), + micro_price_impact: PriceImpact::default(), + micro_variance_ratio: VarianceRatio::default(), + last_timestamp_ns: 0, + last_close: 0.0, + + // WAVE 1.1: Triple barrier integration + triple_barrier, + active_position_tracker: None, + previous_simulated_position: 0.0, // WAVE P3: Start with flat position + + // WAVE 1.2: Safety Infrastructure Integration (8 Systems) + safety_loss_history: VecDeque::with_capacity(30), + safety_loss_plateau_counter: 0, + safety_action_counts: std::collections::HashMap::new(), + safety_memory_manager: Arc::new(RwLock::new( + crate::safety::memory_manager::SafeMemoryManager::new( + &crate::safety::MLSafetyConfig::default() + ) + )), + safety_level: crate::safety::SafetyLevel::Normal, // Default to Normal mode + safety_step_counter: 0, + + feature_cache_dir: None, + + prev_epoch_q_mean: 0.0, + adaptive_tau: base_tau, + + // WAVE 24 (Agent 17): Initialize patience-based early stopping + // Use gradient_collapse_patience from hyperparams for consistency + // Set min_delta to 0.001 (0.1% improvement threshold) + early_stopping: super::super::early_stopping::EarlyStopping::new( + early_stopping_patience, // Reuse patience parameter (default: 5) + 0.001, // 0.1% minimum improvement + ), + + // WAVE 26 P0.6: Use pre-initialized learning rate scheduler + lr_scheduler, + + // WAVE 26 P1.8: Use pre-initialized curiosity module + curiosity_module, + + // WAVE 26 P1: Advanced DQN Features + // P1.3: Sharpe ratio reward + returns_history: VecDeque::with_capacity(sharpe_window), + sharpe_weight, + + // P1.6: Adaptive dropout + dropout_scheduler, + + // P1.7: Hindsight Experience Replay + her_buffer, + + // P1.9: Generalized Advantage Estimation + gae_calculator, + + // P1.11: Noisy sigma scheduler + noisy_sigma_scheduler, + + // WAVE 30: Structured logging integration + logging_config: LoggingConfig::default(), + metrics_aggregator: MetricsAggregator::new(), + + // WAVE 44: Multi-step returns + nstep_buffer, + + // OOM recovery: track effective batch size + current_batch_size: initial_batch_size, + + // Q-value estimation: periodic (every 50 steps) instead of every step + cached_avg_q: 0.0, + q_estimation_counter: 0, + + // GPU pipeline: pre-uploaded training data (initialized lazily at first epoch) + gpu_data: None, + gpu_portfolio_sim: None, + targets_raw_cuda: None, + features_raw_cuda: None, + gpu_experience_collector: None, + gpu_action_selector: None, + training_guard: None, + gpu_monitoring: None, + cached_n_episodes: None, + + // GPU pipeline: staging buffer pool (auto-initialized on CUDA devices) + buffer_pool, + + // GPU pipeline: double-buffered loader for fold transitions + double_buffer, + + // GPU walk-forward: initialized lazily when enable_gpu_walk_forward=true + gpu_walk_forward: None, + + // Multi-GPU: auto-detected data parallelism + multi_gpu, + + // OFI features: populated during data loading when MBP-10 data is available + ofi_features: None, + ofi_val_offset: 0, + val_features_gpu: None, + val_closes_gpu: None, + val_ofi_gpu: None, + + // Phase C: Fill simulation and smart order routing + fill_simulator: FillSimulator::default(), + vol_ema: 0.01, // Initial volatility estimate (1% daily) + median_vol: 0.01, // Slowly adapting baseline + + // Fused CUDA training: lazy-initialized on first training step + fused_ctx: None, + }) + } + +} diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs new file mode 100644 index 000000000..d4d29fd1a --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -0,0 +1,782 @@ +//! DQN Trainer — Training metrics, Q-value diagnostics, and validation + +use anyhow::Result; +use candle_core::Tensor; +use super::DQNTrainer; +use crate::dqn::TradingState; +use crate::dqn::mixed_precision::training_dtype; +use crate::TrainingMetrics; +use super::super::config::DQNAgentType; +use super::super::statistics::QValueStats; + +impl DQNTrainer { + /// Full training loop with existing logic (Wave 12 Group 3) + /// Calculate average metrics for an epoch + pub(crate) fn calculate_epoch_metrics( + epoch_loss: f64, + epoch_q_value: f64, + epoch_gradient_norm: f64, + samples_processed: usize, + ) -> (f64, f64, f64) { + if samples_processed > 0 { + let count = samples_processed as f64; + ( + epoch_loss / count, + epoch_q_value / count, + epoch_gradient_norm / count, + ) + } else { + (0.0, 0.0, 0.0) + } + } + + /// Compute validation loss on held-out data + /// WAVE 10.6: Batched validation for 5-10x speedup + + /// Collect Q-value statistics from replay buffer + /// + /// Samples experiences from the replay buffer and computes Q-value statistics + /// (min, max, mean, std) for adaptive C51 bounds calculation. + /// + /// # Returns + /// + /// QValueStats with min/max/mean/std of Q-values + pub(crate) async fn collect_qvalue_statistics(&self) -> Result { + let agent = self.agent.read().await; + + // Determine sample size (min of buffer size or 1000) + let buffer_size = agent.get_replay_buffer_size()?; + let sample_size = buffer_size.min(1000); + + if sample_size == 0 { + return Err(crate::MLError::TrainingError( + "Replay buffer is empty, cannot collect Q-value statistics".to_owned() + )); + } + + // Sample experiences from replay buffer + let batch_sample = agent.memory().sample(sample_size)?; + + // GPU PER path: use gpu_batch.states directly (always active in CUDA builds) + let gpu_batch = batch_sample.gpu_batch.as_ref() + .ok_or_else(|| crate::MLError::TrainingError( + "GPU PER must be active — gpu_batch is None".to_owned() + ))?; + let batch_tensor = gpu_batch.states.to_dtype(training_dtype(agent.device())) + .map_err(|e| crate::MLError::ModelError(format!("GPU Q-stat states dtype cast: {}", e)))?; + + // Forward pass to get Q-values [batch_size, num_actions] + let q_values = agent.forward(&batch_tensor)?; + + // GPU-side statistics: flatten Q-values and compute min/max/mean/std on device. + // Only 4 scalar readbacks (16 bytes) instead of downloading the entire tensor. + let q_f32 = q_values + .to_dtype(candle_core::DType::F32) + .map_err(|e| crate::MLError::ModelError(format!("Q-value F32 cast: {}", e)))?; + let q_flat = q_f32 + .flatten_all() + .map_err(|e| crate::MLError::ModelError(format!("Q-value flatten: {}", e)))?; + let count = q_flat.elem_count(); + + // Compute all stats on GPU, single batched readback (4 floats in 1 DMA) + let min_t = q_flat.min(0) + .map_err(|e| crate::MLError::ModelError(format!("Q-value min: {}", e)))?; + let max_t = q_flat.max(0) + .map_err(|e| crate::MLError::ModelError(format!("Q-value max: {}", e)))?; + let mean_t = q_flat.mean_all() + .map_err(|e| crate::MLError::ModelError(format!("Q-value mean: {}", e)))?; + let var_t = q_flat.broadcast_sub(&mean_t) + .and_then(|d| d.sqr()) + .and_then(|sq| sq.mean_all()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value variance: {}", e)))?; + + let stats = Tensor::cat( + &[&min_t.unsqueeze(0)?, &max_t.unsqueeze(0)?, &mean_t.unsqueeze(0)?, &var_t.unsqueeze(0)?], 0 + ).and_then(|t| t.to_dtype(candle_core::DType::F32)) + .and_then(|t| t.to_vec1::()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value stats readback: {}", e)))?; + + Ok(QValueStats { + min: *stats.first().unwrap_or(&0.0) as f64, + max: *stats.get(1).unwrap_or(&0.0) as f64, + mean: *stats.get(2).unwrap_or(&0.0) as f64, + std: (*stats.get(3).unwrap_or(&0.0) as f64).sqrt(), + sample_count: count, + }) + } + + /// Create final training metrics + pub(crate) async fn create_final_metrics( + &self, + total_loss: f64, + total_q_value: f64, + total_gradient_norm: f64, + total_reward: f64, + num_epochs: usize, + training_duration: std::time::Duration, + early_stopped: bool, + total_action_counts: [usize; 5], // 5 exposure levels + total_factored_action_counts: [usize; 45], // 45 factored actions + ) -> Result { + let final_loss = total_loss / num_epochs as f64; + let avg_q_value_final = total_q_value / num_epochs as f64; + let avg_grad_norm_final = total_gradient_norm / num_epochs as f64; + let avg_episode_reward = total_reward / num_epochs as f64; + + let mut metrics = TrainingMetrics { + loss: final_loss, + accuracy: 0.0, + precision: 0.0, + recall: 0.0, + f1_score: 0.0, + training_time_seconds: training_duration.as_secs_f64(), + epochs_trained: num_epochs as u32, + convergence_achieved: final_loss < 1.0, + additional_metrics: std::collections::HashMap::new(), + }; + + metrics.add_metric("avg_q_value", avg_q_value_final); + metrics.add_metric("avg_gradient_norm", avg_grad_norm_final); + metrics.add_metric("final_epsilon", self.get_epsilon().await.unwrap_or(0.1)); + metrics.add_metric("avg_episode_reward", avg_episode_reward); + + // Action metrics: use factored 45-action space if branching, else 5 exposure levels + let total_factored: usize = total_factored_action_counts.iter().sum(); + let total_exposure: usize = total_action_counts.iter().sum(); + let total_actions = total_factored.max(total_exposure); + if total_actions > 0 { + // Factored 45-action diversity (primary metric when branching) + if total_factored > 0 { + let unique_factored = total_factored_action_counts.iter() + .filter(|&&count| count > 0).count(); + let factored_diversity = (unique_factored as f64 / 45.0) * 100.0; + metrics.add_metric("action_diversity", factored_diversity); + metrics.add_metric("factored_unique_actions", unique_factored as f64); + metrics.add_metric("action_space_size", 45.0); + + // Active factored actions (used >0.5% of the time) + let active_threshold = (total_factored as f64 * 0.005).max(1.0); + let active_count = total_factored_action_counts.iter() + .filter(|&&count| count as f64 >= active_threshold) + .count(); + let active_diversity_pct = (active_count as f64 / 45.0) * 100.0; + metrics.add_metric("active_actions_count", active_count as f64); + metrics.add_metric("active_diversity_pct", active_diversity_pct); + + // Top factored actions + let mut sorted_actions: Vec<(usize, usize)> = total_factored_action_counts.iter() + .enumerate() + .map(|(idx, &count)| (idx, count)) + .collect(); + sorted_actions.sort_by(|a, b| b.1.cmp(&a.1)); + + if let Some((top1_idx, top1_count)) = sorted_actions.first() { + let top1_pct = (*top1_count as f64 / total_factored as f64) * 100.0; + metrics.add_metric("top1_action_idx", *top1_idx as f64); + metrics.add_metric("top1_action_count", *top1_count as f64); + metrics.add_metric("top1_action_pct", top1_pct); + } + + let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum(); + let top5_coverage_pct = (top5_count as f64 / total_factored as f64) * 100.0; + metrics.add_metric("top5_coverage_pct", top5_coverage_pct); + } else { + // Fallback: 5 exposure levels + let unique_actions = total_action_counts.iter() + .filter(|&&count| count > 0).count(); + let action_diversity = (unique_actions as f64 / 5.0) * 100.0; + metrics.add_metric("action_diversity", action_diversity); + metrics.add_metric("action_space_size", 5.0); + + let active_threshold = (total_exposure as f64 * 0.005).max(1.0); + let active_count = total_action_counts.iter() + .filter(|&&count| count as f64 >= active_threshold) + .count(); + let active_diversity_pct = (active_count as f64 / 5.0) * 100.0; + metrics.add_metric("active_actions_count", active_count as f64); + metrics.add_metric("active_diversity_pct", active_diversity_pct); + + let mut sorted_actions: Vec<(usize, usize)> = total_action_counts.iter() + .enumerate() + .map(|(idx, &count)| (idx, count)) + .collect(); + sorted_actions.sort_by(|a, b| b.1.cmp(&a.1)); + + if let Some((top1_idx, top1_count)) = sorted_actions.first() { + let top1_pct = (*top1_count as f64 / total_exposure as f64) * 100.0; + metrics.add_metric("top1_action_idx", *top1_idx as f64); + metrics.add_metric("top1_action_count", *top1_count as f64); + metrics.add_metric("top1_action_pct", top1_pct); + } + + let top5_count: usize = sorted_actions.iter().take(5).map(|(_, c)| c).sum(); + let top5_coverage_pct = (top5_count as f64 / total_exposure as f64) * 100.0; + metrics.add_metric("top5_coverage_pct", top5_coverage_pct); + } + + metrics.add_metric("total_actions", total_actions as f64); + + // Buy/sell/hold always from 5-exposure space (meaningful for P&L) + // 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100 + let sell_count: usize = total_action_counts.get(0).copied().unwrap_or(0) + + total_action_counts.get(1).copied().unwrap_or(0); + let hold_count: usize = total_action_counts.get(2).copied().unwrap_or(0); + let buy_count: usize = total_action_counts.get(3).copied().unwrap_or(0) + + total_action_counts.get(4).copied().unwrap_or(0); + metrics.add_metric("buy_count", buy_count as f64); + metrics.add_metric("sell_count", sell_count as f64); + metrics.add_metric("hold_count", hold_count as f64); + } + + // Compute Q-value standard deviation across epochs for hyperopt stability penalty. + // self.q_value_history stores per-epoch average Q-values; their std measures + // how much Q-values fluctuate during training (volatility indicator). + if self.q_value_history.len() >= 2 { + let n = self.q_value_history.len() as f64; + let mean = self.q_value_history.iter().sum::() / n; + let variance = self + .q_value_history + .iter() + .map(|&q| (q - mean).powi(2)) + .sum::() + / n; + let std_dev = variance.sqrt(); + metrics.add_metric("q_value_std", std_dev); + } else { + // Not enough data points to compute std; default to 0.0 (no volatility signal) + metrics.add_metric("q_value_std", 0.0); + } + + if early_stopped { + metrics.add_metric("early_stopped", 1.0); + } + + Ok(metrics) + } + + /// Epoch-end Q-value diagnostics: gap analysis + per-action averages. + /// + /// Merges the former `compute_q_gap_for_epoch` and `compute_per_action_q_values` + /// into a single forward pass + readback, eliminating one redundant buffer sample, + /// forward pass, and `to_vec2` GPU-CPU transfer per epoch. + /// + /// Returns (gap_stats, per_action_avgs) where: + /// - gap_stats: (mean_gap, min_gap, max_gap) of Q_best - Q_second_best + /// - per_action_avgs: `[f64; 5]` averages (one per exposure action) + pub(crate) async fn compute_epoch_q_diagnostics(&self) -> Option<( + (f64, f64, f64), + [f64; 5], + )> { + let agent = self.agent.read().await; + let buffer = agent.memory(); + + if buffer.len() < 10 { + return None; + } + + // Use the larger sample size (200) for both diagnostics + let sample_size = buffer.len().min(200); + let batch_sample = match buffer.sample(sample_size) { + Ok(s) => s, + Err(_) => return None, + }; + + let state_dim = agent.get_state_dim(); + + // GPU PER path: use gpu_batch.states directly (experiences vec is empty) + #[allow(unused_mut)] + let mut batch_tensor_opt: Option = None; + { + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = gpu.states + .to_dtype(training_dtype(agent.device())) + .ok(); + } + } + let batch_tensor = if let Some(t) = batch_tensor_opt { + t + } else { + let batched_states: Vec = batch_sample + .experiences + .iter() + .flat_map(|exp| { + let mut s = exp.state.clone(); + s.resize(state_dim, 0.0); + s + }) + .collect(); + let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) { Ok(t) => t, + Err(_) => return None, + }; + match t.to_dtype(training_dtype(agent.device())) { + Ok(t) => t, + Err(_) => return None, + } + }; + + // Single forward pass for both gap and per-action diagnostics + let batch_q_values = match agent.forward(&batch_tensor) { + Ok(q) => q, + Err(_) => return None, + }; + + // All paths use GPU-native diagnostics — no to_vec2 readback + if self.device.is_cuda() { + return Self::compute_q_diagnostics_gpu(&batch_q_values).ok(); + } + + // Non-CUDA: compute diagnostics via tensor ops (no to_vec2) + // Sort Q-values descending per row, gap = sorted[0] - sorted[1] + let sorted = match batch_q_values.sort_last_dim(false) { + Ok((s, _)) => s, + Err(_) => return None, + }; + let n_actions = batch_q_values.dims().get(1).copied().unwrap_or(5); + if n_actions < 2 { return None; } + + let best = match sorted.narrow(1, 0, 1) { + Ok(t) => t.flatten_all().unwrap_or(sorted.clone()), + Err(_) => return None, + }; + let second_best = match sorted.narrow(1, 1, 1) { + Ok(t) => t.flatten_all().unwrap_or(sorted.clone()), + Err(_) => return None, + }; + let gaps_tensor = match best.sub(&second_best) { + Ok(t) => t, + Err(_) => return None, + }; + + let gap_mean = gaps_tensor.mean_all().ok()?.to_scalar::().ok()? as f64; + let gap_min = gaps_tensor.min(0).ok()?.to_scalar::().ok()? as f64; + let gap_max = gaps_tensor.max(0).ok()?.to_scalar::().ok()? as f64; + + // Per-action average Q-values via mean(dim=0) + let per_action = match batch_q_values.mean(0) { + Ok(t) => t, + Err(_) => return None, + }; + + let mut avgs = [0.0_f64; 5]; + for i in 0..5_usize.min(n_actions) { + if let Ok(v) = per_action.narrow(0, i, 1).and_then(|t| t.to_scalar::()) { + avgs[i] = v as f64; + } + } + + Some(((gap_mean, gap_min, gap_max), avgs)) + } + +/// Compute Q-value gap and per-action averages on GPU. +/// Returns (mean_gap, min_gap, max_gap, per_action_avgs[5]). +/// Single 8-float readback at epoch end. +fn compute_q_diagnostics_gpu( + q_values: &Tensor, // [batch, 5] +) -> candle_core::Result<((f64, f64, f64), [f64; 5])> { + // sort_last_dim returns (sorted_values, indices) — destructure the tuple + let (sorted, _indices) = q_values.sort_last_dim(true)?; // descending + let best = sorted.narrow(1, 0, 1)?; + let second = sorted.narrow(1, 1, 1)?; + let gaps = best.sub(&second)?; + + // Compute gap stats on GPU, batch into one tensor for single readback + let gaps_flat = gaps.flatten_all()?; + let mean_gap = gaps_flat.mean_all()?; + let min_gap = gaps_flat.min(0)?; + let max_gap = gaps_flat.max(0)?; + + // Per-action means: mean along batch dim [5] + let per_action = q_values.mean(0)?; + + // Single readback: cat [mean, min, max, per_action_0..4] → 8 floats in one DMA + let gap_vec = Tensor::cat( + &[&mean_gap.unsqueeze(0)?, &min_gap.unsqueeze(0)?, &max_gap.unsqueeze(0)?], 0 + )?; + let all_stats = Tensor::cat(&[&gap_vec, &per_action.flatten_all()?], 0)? + .to_dtype(candle_core::DType::F32)? + .to_vec1::()?; + + let mean_g = *all_stats.first().unwrap_or(&0.0) as f64; + let min_g = *all_stats.get(1).unwrap_or(&0.0) as f64; + let max_g = *all_stats.get(2).unwrap_or(&0.0) as f64; + let mut avgs = [0.0_f64; 5]; + for i in 0..5_usize { + avgs[i] = *all_stats.get(3 + i).unwrap_or(&0.0) as f64; + } + + Ok(((mean_g, min_g, max_g), avgs)) +} + + /// Get current training metrics + pub async fn get_metrics(&self) -> TrainingMetrics { + self.metrics.read().await.clone() + } + + /// Get Q-values for a given state + pub(crate) async fn get_q_values(&self, state: &TradingState) -> Result> { + let agent = self.agent.read().await; + let state_vec = state.to_vector(); + let raw_dim = state_vec.len(); + let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); + let padded: Vec = if aligned > raw_dim { + let mut v = state_vec.to_vec(); + v.resize(aligned, 0.0); + v + } else { + state_vec.to_vec() + }; + let state_tensor = Tensor::new(&*padded, &self.device)?.unsqueeze(0)?; // Add batch dimension + + let q_values_tensor = agent.forward(&state_tensor)?.squeeze(0)?; + // Single readback: download entire Q-value vector in one DMA + let q_f32 = q_values_tensor + .to_dtype(candle_core::DType::F32)? + .to_vec1::()?; + Ok(q_f32.into_iter().map(|v| v as f64).collect()) + } + + /// Check if early stopping criteria are met + pub(crate) fn check_early_stopping(&self, avg_q_value: f64, epoch: usize) -> Option { + if !self.hyperparams.early_stopping_enabled + || epoch + 1 < self.hyperparams.min_epochs_before_stopping + { + return None; + } + + // Criterion 1: Q-value floor check + if avg_q_value < self.hyperparams.q_value_floor { + return Some(format!( + "Q-value {:.4} below floor threshold {:.4}", + avg_q_value, self.hyperparams.q_value_floor + )); + } + + // C4 FIX: Criterion 2 — Sharpe plateau check (was val-loss plateau). + // Sharpe directly measures trading quality. A plateau means the model + // has stopped improving its trading strategy, even if TD-loss still moves. + if self.sharpe_history.len() >= self.hyperparams.plateau_window { + let window = self.hyperparams.plateau_window; + let recent_sharpes: Vec = self + .sharpe_history + .iter() + .rev() + .take(window) + .copied() + .collect(); + + if let (Some(&newest), Some(&oldest)) = (recent_sharpes.first(), recent_sharpes.last()) { + // For Sharpe, improvement = newest - oldest (higher is better) + let improvement = newest - oldest; + + if improvement < 0.01 { + let msg = if improvement < -0.01 { + format!( + "Sharpe worsening detected (delta: {:.4}, window: {})", + improvement, + window + ) + } else { + format!( + "Sharpe plateau detected (improvement: {:.4}, window: {})", + improvement, + window + ) + }; + return Some(msg); + } + } + } + + None + } + + #[allow(unused_variables, unreachable_code, unused_mut)] + pub(crate) async fn compute_validation_loss(&mut self) -> Result { + if self.val_data.is_empty() { + return Ok(0.0); + } + + // Save current epsilon and force to 0 for deterministic evaluation + let original_epsilon = self.get_epsilon().await?; + self.set_epsilon(0.0).await?; // Pure greedy selection + + let sample_size = self.val_data.len().min(1000); // Sample up to 1000 for speed + + let aligned_state_dim = { + let agent = self.agent.read().await; + agent.get_state_dim() + }; + + // ── GPU-resident validation data (lazy init, uploaded once) ── + // Pre-upload val features [sample_size, 42], close prices, and OFI features to GPU. + // Subsequent epochs reuse the same GPU tensors — zero CPU loop. + if self.val_features_gpu.is_none() { + let mut flat_features = Vec::with_capacity(sample_size * 42); + let mut current_closes = Vec::with_capacity(sample_size); + let mut next_closes = Vec::with_capacity(sample_size); + + for (feature_vec, target) in self.val_data.iter().take(sample_size) { + for &v in feature_vec.iter() { + flat_features.push(v as f32); + } + let cur = if target.len() >= 2 { target[0] } else { feature_vec[3] }; + let nxt = if target.len() >= 2 { target[1] } else { cur }; + current_closes.push(cur as f32); + next_closes.push(nxt as f32); + } + + self.val_features_gpu = Some( + Tensor::from_vec(flat_features, (sample_size, 42), &self.device) + .map_err(|e| anyhow::anyhow!("GPU val features upload: {e}"))? + ); + let cur_t = Tensor::from_vec(current_closes, &[sample_size], &self.device) + .map_err(|e| anyhow::anyhow!("GPU val current_closes upload: {e}"))?; + let nxt_t = Tensor::from_vec(next_closes, &[sample_size], &self.device) + .map_err(|e| anyhow::anyhow!("GPU val next_closes upload: {e}"))?; + self.val_closes_gpu = Some((cur_t, nxt_t)); + + // Upload OFI features for validation range if available + if let Some(ref ofi) = self.ofi_features { + let mut flat_ofi = Vec::with_capacity(sample_size * 8); + for i in 0..sample_size { + let idx = self.ofi_val_offset + i; + if let Some(row) = ofi.get(idx) { + for &v in row.iter() { + flat_ofi.push(v as f32); + } + } else { + flat_ofi.extend_from_slice(&[0.0_f32; 8]); + } + } + self.val_ofi_gpu = Some( + Tensor::from_vec(flat_ofi, (sample_size, 8), &self.device) + .map_err(|e| anyhow::anyhow!("GPU val OFI upload: {e}"))? + ); + } + } + + // ── Build state tensor on GPU: cat [features, portfolio, ofi] ── + let features_gpu = self.val_features_gpu.as_ref() + .ok_or_else(|| anyhow::anyhow!("val_features_gpu must be initialized"))?; + + // Portfolio features: 3 scalars from current tracker, broadcast to [sample_size, 3] + // Use first validation sample's close price for portfolio evaluation + let val_price = self.val_data.first() + .map(|(fv, tgt)| if tgt.len() >= 2 { tgt[0] as f32 } else { fv[3] as f32 }) + .unwrap_or(0.0); + let portfolio_f = self.portfolio_tracker.get_portfolio_features(val_price); + let portfolio_gpu = Tensor::new(&portfolio_f[..], &self.device) + .map_err(|e| anyhow::anyhow!("GPU portfolio tensor: {e}"))? + .unsqueeze(0) + .map_err(|e| anyhow::anyhow!("GPU portfolio unsqueeze: {e}"))? + .broadcast_left(sample_size) + .map_err(|e| anyhow::anyhow!("GPU portfolio broadcast: {e}"))?; + + // Concat features + portfolio (+ OFI if enabled) → [sample_size, raw_state_dim] + let raw_state_dim = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; + let state_gpu = if let Some(ref ofi_gpu) = self.val_ofi_gpu { + Tensor::cat(&[features_gpu, &portfolio_gpu, ofi_gpu], 1) + .map_err(|e| anyhow::anyhow!("GPU val state cat (with OFI): {e}"))? + } else { + Tensor::cat(&[features_gpu, &portfolio_gpu], 1) + .map_err(|e| anyhow::anyhow!("GPU val state cat: {e}"))? + }; + + // Pad to aligned dim if needed (trailing zeros for tensor core alignment) + let batch_tensor = if aligned_state_dim > raw_state_dim { + let pad_width = aligned_state_dim - raw_state_dim; + let pad = Tensor::zeros((sample_size, pad_width), candle_core::DType::F32, &self.device) + .map_err(|e| anyhow::anyhow!("GPU val pad zeros: {e}"))?; + Tensor::cat(&[&state_gpu, &pad], 1) + .map_err(|e| anyhow::anyhow!("GPU val state pad: {e}"))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("GPU val state dtype: {e}"))? + } else { + state_gpu + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("GPU val state dtype: {e}"))? + }; + + // Close prices already on GPU from lazy init + let (ref val_current_closes_t, ref val_next_closes_t) = self.val_closes_gpu.as_ref() + .ok_or_else(|| anyhow::anyhow!("val_closes_gpu must be initialized"))?; + + let agent = self.agent.read().await; + + let batch_q_values = agent.forward(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Batched validation forward pass failed: {}", e))?; + + // Get branching Q-values if branching is enabled + let branching_q_tensors: Option<(Tensor, Tensor, Tensor)> = + if self.hyperparams.use_branching { + agent + .batch_branching_q_values(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Validation branching Q-values failed: {e}"))? + } else { + None + }; + + drop(agent); // Release lock early + + + // Fused GPU greedy action selection (epsilon=0.0) + GPU routing. + { + if self.gpu_action_selector.is_none() && self.device.is_cuda() { + let selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new( + &self.device, + self.hyperparams.batch_size.max(sample_size).max(8192), + 0xDEAD_BEEF_CAFE_u64, + ).map_err(|e| anyhow::anyhow!("Validation GPU action selector init failed: {e}"))?; + self.gpu_action_selector = Some(selector); + } + + let selector = self.gpu_action_selector.as_mut() + .ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?; + + let factored_tensor = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors { + selector + .select_actions_branching(q_exp, q_ord, q_urg, 0.0) + .map_err(|e| anyhow::anyhow!("Validation GPU branching select failed: {e}"))? + } else { + let exposure_tensor = selector + .select_actions(&batch_q_values, 0.0, sample_size, 5) + .map_err(|e| anyhow::anyhow!("Validation GPU greedy select failed: {e}"))?; + selector.route_exposure_to_factored( + &exposure_tensor, sample_size, + self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, + self.vol_ema as f32, self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("Validation GPU route failed: {e}"))? + }; + + // --- GPU PnL-based reward + Sharpe computation --- + // Instead of reading factored indices back to CPU and looping through + // the scalar reward function, compute PnL rewards entirely on GPU: + // reward_i = direction_i * (next_close_i - current_close_i) / current_close_i + // Then compute Sharpe = mean(rewards) / std(rewards) * sqrt(252) on device + // with a single scalar readback for the final value. + + // 1) Close prices already on GPU (pre-uploaded) + let current_closes_t: &Tensor = val_current_closes_t; + let next_closes_t: &Tensor = val_next_closes_t; + + // 2) Extract exposure index from factored tensor: + // factored_index = exposure * 9 + order * 3 + urgency + // => exposure_idx = factored_index / 9 + let factored_f32 = factored_tensor + .to_dtype(candle_core::DType::F32) + .map_err(|e| anyhow::anyhow!("GPU val factored to f32: {e}"))?; + let nine = Tensor::new(&[9.0_f32], &self.device) + .map_err(|e| anyhow::anyhow!("GPU val nine const: {e}"))?; + let exposure_idx_f32 = factored_f32 + .broadcast_div(&nine) + .map_err(|e| anyhow::anyhow!("GPU val exposure div: {e}"))? + .floor() + .map_err(|e| anyhow::anyhow!("GPU val exposure floor: {e}"))?; + + // 3) Map exposure index to direction multiplier via lookup table: + // [0=Short100→-1.0, 1=Short50→-0.5, 2=Flat→0.0, 3=Long50→0.5, 4=Long100→1.0] + let direction_lut = Tensor::new( + &[-1.0_f32, -0.5, 0.0, 0.5, 1.0], &self.device, + ).map_err(|e| anyhow::anyhow!("GPU val direction LUT: {e}"))?; + let exposure_idx_u32 = exposure_idx_f32 + .to_dtype(candle_core::DType::U32) + .map_err(|e| anyhow::anyhow!("GPU val exposure to u32: {e}"))?; + let directions = direction_lut + .index_select(&exposure_idx_u32, 0) + .map_err(|e| anyhow::anyhow!("GPU val direction gather: {e}"))?; + + // 4) Compute PnL-based rewards: direction * (next - current) / current + let price_returns = next_closes_t + .sub(¤t_closes_t) + .map_err(|e| anyhow::anyhow!("GPU val price diff: {e}"))? + .broadcast_div(¤t_closes_t) + .map_err(|e| anyhow::anyhow!("GPU val price returns div: {e}"))?; + let rewards = directions + .mul(&price_returns) + .map_err(|e| anyhow::anyhow!("GPU val rewards mul: {e}"))?; + + // 5) Compute Sharpe on GPU: mean / std * sqrt(252) — single batched readback + let mean_t = rewards.mean_all() + .map_err(|e| anyhow::anyhow!("GPU val rewards mean: {e}"))?; + let var_t = rewards + .broadcast_sub(&mean_t) + .map_err(|e| anyhow::anyhow!("GPU val rewards center: {e}"))? + .sqr() + .map_err(|e| anyhow::anyhow!("GPU val rewards sqr: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU val rewards var: {e}"))?; + let stats = Tensor::cat(&[&mean_t.unsqueeze(0)?, &var_t.unsqueeze(0)?], 0) + .and_then(|t| t.to_dtype(candle_core::DType::F32)) + .and_then(|t| t.to_vec1::()) + .map_err(|e| anyhow::anyhow!("GPU val Sharpe stats readback: {e}"))?; + let mean_scalar = *stats.first().unwrap_or(&0.0) as f64; + let var_scalar = *stats.get(1).unwrap_or(&0.0) as f64; + + let std_val = var_scalar.sqrt(); + let val_sharpe = if std_val > 1e-10 { + (mean_scalar / std_val) * (252.0_f64).sqrt() + } else { + 0.0 + }; + + // Restore original epsilon after evaluation + self.set_epsilon(original_epsilon).await?; + + // Return negative Sharpe as the "loss" (lower = better Sharpe) + return Ok(-val_sharpe); + } + } + + /// Estimate average Q-value from replay buffer samples for monitoring + /// + /// WAVE 23 P0: Now includes Q-value divergence check (early stopping) + /// OPTIMIZATION: Batched Q-value estimation for 10× speedup via GPU parallelization + pub(crate) async fn estimate_avg_q_value_with_early_stopping(&self, agent: &mut DQNAgentType) -> Result { + // Get a few samples from the replay buffer to estimate Q-values + let buffer = agent.memory(); + + if buffer.len() == 0 { + return Ok(0.0); + } + + // Sample up to 10 experiences for Q-value estimation + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Failed to sample experiences: {}", e))?; + + // GPU PER path: use gpu_batch.states directly (always active in CUDA builds) + let gpu_batch = batch_sample.gpu_batch.as_ref() + .ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?; + let batch_tensor = gpu_batch.states.to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("GPU Q-est states dtype cast: {}", e))?; + + // WAVE 23 P0 Fix: Check for Q-value divergence (early stopping) + // This calls log_q_values() which returns Err if divergence detected for consecutive checks + agent.log_q_values(&batch_tensor) + .map_err(|e| { + tracing::info!("🛑 Early stopping triggered (Q-value divergence): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Single forward pass for all samples (10× faster than sequential) + let batch_q_values = agent + .forward(&batch_tensor) + .map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?; + + // Get max Q-value per sample across action dimension + let max_q_values = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("Failed to compute max Q-values: {}", e))?; + + // Compute average across batch + let avg_q = max_q_values + .mean_all() + .map_err(|e| anyhow::anyhow!("Failed to compute mean Q-value: {}", e))? + .to_scalar::() .map_err(|e| anyhow::anyhow!("Failed to extract average Q-value: {}", e))? + as f64; + + Ok(avg_q) + } + +} diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs new file mode 100644 index 000000000..280beaee3 --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -0,0 +1,872 @@ +//! DQN Trainer Implementation +//! +//! Main training loop and execution logic for Deep Q-Network. + +use std::collections::VecDeque; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use candle_core::{Device, Tensor}; +use crate::cuda_pipeline::DqnGpuData; +use ml_core::fill_simulator::FillSimulator; +use risk::drawdown_monitor::DrawdownMonitor; +use risk::safety::position_limiter::HybridPositionLimiter; +use tokio::sync::RwLock; +use tracing::info; +use uuid::Uuid; + +use crate::dqn::action_space::FactoredAction; +use crate::dqn::circuit_breaker::CircuitBreaker; +use crate::dqn::curiosity::CuriosityModule; +use crate::dqn::logging::{LoggingConfig, MetricsAggregator}; +use crate::dqn::portfolio_tracker::PortfolioTracker; +use crate::dqn::reward::RewardFunction; +use crate::TrainingMetrics; +use crate::features::extraction::FeatureVector; +use crate::features::microstructure_features::*; +use crate::labeling::triple_barrier::TripleBarrierEngine; + +use super::config::{DQNAgentType, DQNHyperparameters}; +use super::statistics::{FeatureStatistics, QValueStats}; +pub(super) use super::EPISODE_LENGTH; + +mod action; +mod metrics; +mod state; +mod constructor; +mod training_loop; +mod train_step; + +#[cfg(test)] +mod tests; + + + +pub struct DQNTrainer { + /// DQN agent + pub(crate) agent: Arc>, + /// Training hyperparameters + pub(crate) hyperparams: DQNHyperparameters, + /// Device (GPU or CPU) + pub(crate) device: Device, + /// Training metrics + pub(crate) metrics: Arc>, + /// Loss history for plateau detection + pub(crate) loss_history: Vec, + /// Q-value history for floor detection + pub(crate) q_value_history: Vec, + /// Best validation loss achieved so far + pub(crate) best_val_loss: f64, + /// Validation data for computing validation loss + pub(crate) val_data: Vec<(FeatureVector, Vec)>, + /// Validation loss history for early stopping + pub(crate) val_loss_history: Vec, + /// C4: Sharpe history for Sharpe-based early stopping + pub(crate) sharpe_history: Vec, + /// C4: Best Sharpe ratio achieved so far (higher = better) + pub(crate) best_sharpe: f64, + /// Epoch with best validation loss + pub(crate) best_epoch: usize, + /// Step counter for gradient logging (logs every 10 steps) + pub(crate) gradient_logging_step: usize, + /// Original buffer_size before AutoReplaySizer (for gradient collapse warmup) + pub(crate) collapse_warmup_buffer_size: usize, + /// Portfolio state tracker for P&L-based rewards (Bug #2 fix) + pub portfolio_tracker: PortfolioTracker, + /// Feature normalization statistics (WAVE 3 FIX #2) + /// None during stats collection phase (epochs 0-10), Some during normalization phase (epochs 11+) + pub feature_stats: Option, + /// Sliding window of recent actions for reward calculation (max 100) + pub(crate) recent_actions: VecDeque, + /// Reward function for calculating rewards with recent actions + pub(crate) reward_fn: RewardFunction, + + // WAVE 16S: Adaptive Risk Management Components + /// Kelly criterion optimizer for position sizing (None if disabled) + pub(crate) kelly_optimizer: Option>, + /// Trade history for Kelly calculation (wins/losses) + pub(crate) trade_history: VecDeque, + /// Volatility tracker for epsilon adjustment (None if disabled) + pub(crate) volatility_returns: VecDeque, + /// PnL history for Sharpe calculation (max 1000 entries) + pub(crate) pnl_history: VecDeque, + + // Wave 16 Portfolio Features + /// Enable action masking (filters invalid actions before Q-value computation) + pub enable_action_masking: bool, + /// Maximum position size for action masking (default: 2.0) + pub max_position: f64, + /// Entropy regularizer for preventing policy collapse (None if disabled) + // entropy_regularizer removed — SAC-style entropy is computed directly on Q-value tensors in DQN::compute_loss_internal + /// Multi-asset portfolio tracker (None if single-asset mode) + pub multi_asset_portfolio: Option>, + /// Stress tester for robustness validation (None if disabled) + pub stress_tester: Option>, + + // Wave 16 Core Risk Features Integration + /// Drawdown monitor for tracking portfolio drawdowns (15% max drawdown) + pub drawdown_monitor: Option>, + /// Position limiter with 3-tier limits (±10.0 absolute, 1M notional, 10% concentration) + pub position_limiter: Option>, + /// Circuit breaker for stopping training on consecutive failures + pub circuit_breaker: Option>, + + // WAVE 3.10: Microstructure Feature Calculators (12 features) + pub(crate) micro_high_low_spread: HighLowSpread, + pub(crate) micro_vw_spread: VolumeWeightedSpread, + pub(crate) micro_tick_count: TickCount, + pub(crate) micro_inter_arrival: InterArrivalTime, + pub(crate) micro_buy_sell_imbalance: BuySellImbalance, + pub(crate) micro_kyle_lambda: KyleLambda, + pub(crate) micro_price_impact: PriceImpact, + pub(crate) micro_variance_ratio: VarianceRatio, + // Note: Roll Measure, Corwin-Schultz, Amihud, VPIN already exist in ml/src/microstructure/ + // We'll integrate those in the update logic + /// Track last timestamp for inter-arrival time calculation + pub(crate) last_timestamp_ns: u64, + /// Track last close price for microstructure calculations + pub(crate) last_close: f64, + + // WAVE 1.1: Triple Barrier Integration + /// Triple barrier engine for position exit labeling + pub(crate) triple_barrier: Arc>, + /// Active position tracker ID (None = no active position) + pub(crate) active_position_tracker: Option, + /// WAVE P3: Track previous simulated position for barrier tracking continuity + pub(crate) previous_simulated_position: f32, + + // WAVE 1.2: Safety Infrastructure Integration (8 Systems) + /// Loss history window for spike detection (size: 30) + pub(crate) safety_loss_history: VecDeque, + /// Loss plateau counter for anomaly detection + pub(crate) safety_loss_plateau_counter: usize, + /// Action counts for diversity monitoring (5 exposure actions) + pub(crate) safety_action_counts: std::collections::HashMap, + /// Memory manager for GPU OOM risk monitoring + pub(crate) safety_memory_manager: Arc>, + /// Safety enforcement level (Strict/Normal/Permissive) + pub(crate) safety_level: crate::safety::SafetyLevel, + /// Step counter for periodic safety checks + pub(crate) safety_step_counter: usize, + + /// Optional path to feature cache directory for faster hyperopt + pub(crate) feature_cache_dir: Option, + + /// Previous epoch's mean Q-value for overestimation detection + pub(crate) prev_epoch_q_mean: f64, + /// Current effective tau (may be temporarily increased if Q-values drift) + pub(crate) adaptive_tau: f64, + + /// WAVE 24 (Agent 17): Patience-based early stopping for anti-overfitting + pub(crate) early_stopping: super::early_stopping::EarlyStopping, + + /// WAVE 26 P0.6: Learning rate scheduler with warmup + pub(crate) lr_scheduler: super::lr_scheduler::LRScheduler, + + /// WAVE 26 P1.8: Curiosity module for intrinsic rewards (None if curiosity_weight = 0.0) + pub(crate) curiosity_module: Option, + + // WAVE 26 P1: Advanced DQN Features Integration + // P1.3: Sharpe Ratio Reward Component + /// Rolling buffer of returns for Sharpe ratio calculation (max: sharpe_window) + pub(crate) returns_history: VecDeque, + /// Sharpe reward weight (0.0 = disabled) + pub(crate) sharpe_weight: f64, + + // P1.6: Adaptive Dropout Scheduling + /// Optional dropout scheduler (None if disabled) + pub(crate) dropout_scheduler: Option, + + // P1.7: Hindsight Experience Replay (HER) + /// Optional HER buffer (None if her_ratio = 0.0) + pub(crate) her_buffer: Option>, + + // P1.9: Generalized Advantage Estimation (GAE) + /// Optional GAE calculator (None if disabled) + pub(crate) gae_calculator: Option, + + // P1.11: Noisy Network Sigma Scheduling + /// Optional noisy sigma scheduler (None if disabled) + pub(crate) noisy_sigma_scheduler: Option, + + // WAVE 30: Structured Logging Integration + /// Logging configuration for training metrics + pub(crate) logging_config: LoggingConfig, + /// Metrics aggregator for windowed training statistics + pub(crate) metrics_aggregator: MetricsAggregator, + + // WAVE 44: Multi-step returns integration + /// N-step buffer for multi-step TD learning (None if n_steps=1) + pub(crate) nstep_buffer: Option, + + /// Current effective batch size (may be reduced by OOM recovery) + pub(crate) current_batch_size: usize, + + /// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step) + pub(crate) cached_avg_q: f64, + /// Counter for Q-value estimation frequency (estimate every N training steps) + pub(crate) q_estimation_counter: u64, + + /// Pre-uploaded GPU training data (set once, reused across epochs) + pub(crate) gpu_data: Option, + + /// GPU portfolio simulator for CUDA-accelerated experience collection + pub(crate) gpu_portfolio_sim: Option, + + /// Raw cudarc targets buffer for CUDA kernel (parallel to candle Tensor in gpu_data) + pub(crate) targets_raw_cuda: Option>, + + /// Raw cudarc features buffer for CUDA experience kernel [num_bars * 42] + pub(crate) features_raw_cuda: Option>, + + /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b) + pub(crate) gpu_experience_collector: Option, + + /// GPU-fused epsilon-greedy action selector (eliminates argmax GPU->CPU sync barrier) + pub(crate) gpu_action_selector: Option, + + /// GPU training guard for zero-sync safety checks (loss clip, NaN, grad collapse) + pub(crate) training_guard: Option, + + /// GPU monitoring reducer — accumulates reward/action stats across kernel launches + pub(crate) gpu_monitoring: Option, + + /// Reusable GPU staging buffers for zero-alloc fold transitions + pub(crate) buffer_pool: Option, + + /// Double-buffered GPU data for zero-downtime fold transitions + pub(crate) double_buffer: Option, + + /// GPU-resident walk-forward data (entire dataset on GPU, per-fold views via index ranges) + pub(crate) gpu_walk_forward: Option, + + /// Multi-GPU configuration for data-parallel training (None = single GPU) + pub(crate) multi_gpu: Option, + + /// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs) + /// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork). + pub(crate) cached_n_episodes: Option, + + /// Pre-computed OFI features per bar (indexed by global bar position). + /// Populated during data loading when MBP-10 order book data is available. + /// Passed as `regime_features` in `TradingState::from_normalized()`. + /// Arc-shared to avoid 2.68 GB copy per hyperopt trial (41.9M × 8 × 8 bytes). + pub(crate) ofi_features: Option>, + /// Number of training bars (OFI offset for validation data). + /// val_data[i] corresponds to ofi_features[ofi_val_offset + i]. + pub(crate) ofi_val_offset: usize, + + /// GPU-resident validation features [sample_size, 42] — pre-uploaded once, reused per epoch + pub(crate) val_features_gpu: Option, + /// GPU-resident validation close prices (current, next) — pre-uploaded once + pub(crate) val_closes_gpu: Option<(Tensor, Tensor)>, + /// GPU-resident validation OFI features [sample_size, 8] — pre-uploaded once + pub(crate) val_ofi_gpu: Option, + + // Phase C: Fill simulation and smart order routing + /// Fill simulator for order type-dependent execution modeling + pub(crate) fill_simulator: FillSimulator, + /// EMA of bar volatility (|close log return|) for OrderRouter routing decisions + pub(crate) vol_ema: f64, + /// Running median volatility estimate (slowly adapting EMA) + pub(crate) median_vol: f64, + + /// Fused CUDA training context: pre-allocated batch buffers + single entry point + /// for CUDA Graph capture. Only active for Standard DQN on CUDA devices. + /// Lazy-initialized on first training step; dropped and recreated if batch_size changes. + pub(crate) fused_ctx: Option, +} + +impl std::fmt::Debug for DQNTrainer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DQNTrainer") + .field("hyperparams", &self.hyperparams) + .finish_non_exhaustive() + } +} + +impl DQNTrainer { + /// Create new DQN trainer with hyperparameters and debug logging disabled + pub fn new(hyperparams: DQNHyperparameters) -> Result { + Self::new_with_debug(hyperparams, false) + } + + /// Create new DQN trainer with a specific compute device. + /// Used by hyperopt to share a single CUDA context across parallel trials. + pub fn new_with_device(hyperparams: DQNHyperparameters, device: Device) -> Result { + Self::new_internal(hyperparams, false, Some(device)) + } + + /// Create new DQN trainer with hyperparameters and configurable debug logging + /// + /// # Arguments + /// * `hyperparams` - DQN training hyperparameters + /// * `debug_logging` - Enable debug logging (REWARD_DEBUG, gradient norms, etc.) + pub fn new_with_debug(hyperparams: DQNHyperparameters, debug_logging: bool) -> Result { + Self::new_internal(hyperparams, debug_logging, None) + } + + + /// Two-phase stress tester initialization. + /// + /// Call this after constructing `DQNTrainer` to resolve the circular dependency: + /// `DQNStressTester::new()` requires a `DQNTrainer`, so the tester cannot be + /// created *during* trainer construction. This method builds a lightweight + /// inner trainer (with stress testing itself disabled to avoid recursion) and + /// hands it to `DQNStressTester::new()`. + /// + /// No-op if `hyperparams.enable_stress_testing` is `false`. + pub fn init_stress_tester(&mut self) -> Result<()> { + if !self.hyperparams.enable_stress_testing { + return Ok(()); + } + + // Clone hyperparams with stress testing disabled so the inner trainer + // does not recursively try to initialise its own stress tester. + // Also disable GPU-heavy features that the stress tester doesn't need — + // otherwise we double the VRAM usage (3 extra regime heads, 3 extra GPU PER buffers, + // experience collector, etc.) which causes OOM on 4GB GPUs. + let mut inner_hp = self.hyperparams.clone(); + inner_hp.enable_stress_testing = false; + inner_hp.enable_regime_qnetwork = false; + inner_hp.use_per = false; + inner_hp.buffer_size = 1; + inner_hp.enable_gpu_experience_collector = false; + + let inner_trainer = Self::new_with_device(inner_hp, self.device.clone()) + .context("Failed to create inner DQNTrainer for stress tester")?; + + let tester = crate::dqn::stress_testing::DQNStressTester::new(inner_trainer)?; + self.stress_tester = Some(Arc::new(tester)); + info!("Stress testing enabled (8 scenarios)"); + + Ok(()) + } + + /// Get a reference to the double-buffered loader, if GPU is active. + pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> { + self.double_buffer.as_ref() + } + + /// Get a mutable reference to the double-buffered loader, if GPU is active. + pub fn double_buffer_mut(&mut self) -> Option<&mut crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> { + self.double_buffer.as_mut() + } + + /// Set feature cache directory for faster hyperopt + /// + /// Enables loading pre-computed features from disk instead of recomputing them + pub fn with_feature_cache(mut self, cache_dir: PathBuf) -> Self { + self.feature_cache_dir = Some(cache_dir); + self + } + + /// Train DQN on market data from DBN files + /// + /// # Arguments + /// + /// * `dbn_data_dir` - Directory containing DBN files (e.g., "test_data/real/databento/ml_training/") + /// * `checkpoint_callback` - Callback for saving checkpoints (epoch, model_data, is_final) -> `Result` + /// + /// # Returns + /// + /// Training metrics (loss, accuracy, gradient norms, Q-values) + pub async fn train( + &mut self, + dbn_data_dir: &str, + checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + info!( + "Starting DQN training for {} epochs with batch size {}", + self.hyperparams.epochs, self.hyperparams.batch_size + ); + + // Load market data from DBN files (ALL data for walk-forward or single-pass) + let (training_data, val_data) = self.load_training_data(dbn_data_dir).await?; + + info!( + "Loaded {} training samples, {} validation samples", + training_data.len(), + val_data.len() + ); + + // GPU walk-forward: upload ALL data to GPU, run expanding-window folds + if self.hyperparams.enable_gpu_walk_forward && self.device.is_cuda() { + // Merge train+val into a single dataset for walk-forward splitting + let mut all_data = training_data; + all_data.extend(val_data); + info!( + "GPU walk-forward enabled: {} total bars, uploading to VRAM", + all_data.len(), + ); + return self.train_walk_forward(&all_data, checkpoint_callback).await; + } + + // Standard single-pass training + self.ofi_val_offset = training_data.len(); + self.val_data = val_data; + self.val_features_gpu = None; + self.val_closes_gpu = None; + self.val_ofi_gpu = None; + self.train_with_data_full_loop(&training_data, checkpoint_callback) + .await + } + + /// Train with preloaded data (skips disk I/O and feature extraction). + /// + /// Accepts pre-split training and validation data that was loaded once and + /// cached across hyperopt trials. This avoids re-reading 36 `.dbn.zst` files + /// and re-extracting 42 features on every trial, eliminating minutes of GPU + /// idle time at each trial boundary. + /// + /// # Arguments + /// + /// * `training_data` - Pre-extracted (features, targets) for training split + /// * `val_data` - Pre-extracted (features, targets) for validation split + /// * `checkpoint_callback` - Checkpoint save callback + /// + /// # Returns + /// + /// Training metrics from the completed run + /// + /// # Errors + /// + /// Returns error if the training loop fails + pub async fn train_with_preloaded_data( + &mut self, + training_data: Vec<(FeatureVector, Vec)>, + val_data: Vec<(FeatureVector, Vec)>, + checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + info!( + "Starting DQN training with preloaded data: {} train, {} val samples", + training_data.len(), + val_data.len() + ); + + // Store validation data for loss computation + self.ofi_val_offset = training_data.len(); + self.val_data = val_data; + self.val_features_gpu = None; + self.val_closes_gpu = None; + self.val_ofi_gpu = None; + + // Use the common training loop (Wave 12 Group 3 refactor) + self.train_with_data_full_loop(&training_data, checkpoint_callback) + .await + } + + /// Train with shared preloaded data (zero-copy for hyperopt). + /// + /// Same as [`train_with_preloaded_data`] but accepts `Arc`-wrapped data, + /// avoiding a ~150 MB deep clone per hyperopt trial. + pub async fn train_with_shared_data( + &mut self, + training_data: &[(FeatureVector, Vec)], + val_data: Vec<(FeatureVector, Vec)>, + checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + info!( + "Starting DQN training with shared data: {} train, {} val samples", + training_data.len(), + val_data.len() + ); + + self.ofi_val_offset = training_data.len(); + self.val_data = val_data; + self.val_features_gpu = None; + self.val_closes_gpu = None; + self.val_ofi_gpu = None; + self.train_with_data_full_loop(training_data, checkpoint_callback) + .await + } + + /// Train with GPU-resident walk-forward cross-validation. + /// + /// Uploads the ENTIRE dataset to GPU VRAM once, then runs expanding-window + /// walk-forward: each fold trains on [0..T], validates on [T..V], tests on + /// [V..E]. Fold transitions are zero-copy (index range changes only). + /// + /// Returns the metrics from the LAST fold (most data, most representative). + pub async fn train_walk_forward( + &mut self, + training_data: &[(FeatureVector, Vec)], + mut checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + use crate::cuda_pipeline::gpu_walk_forward::{GpuWalkForwardConfig, GpuWalkForwardData}; + + let wf_config = GpuWalkForwardConfig { + initial_train_fraction: self.hyperparams.wf_initial_train_fraction, + val_fraction: self.hyperparams.wf_val_fraction, + test_fraction: self.hyperparams.wf_test_fraction, + step_fraction: self.hyperparams.wf_step_fraction, + }; + + // Upload ALL data to GPU once + let wf_data = GpuWalkForwardData::upload( + training_data, + self.ofi_features.as_deref(), + &wf_config, + &self.device, + ).map_err(|e| anyhow::anyhow!("GPU walk-forward upload: {e}"))?; + + let num_folds = wf_data.num_folds(); + if num_folds == 0 { + return Err(anyhow::anyhow!( + "Insufficient data for walk-forward: {} bars, need at least {} for one fold", + training_data.len(), + ((wf_config.initial_train_fraction + wf_config.val_fraction + wf_config.test_fraction) * training_data.len() as f64) as usize, + )); + } + + info!( + "GPU walk-forward: {} folds, {:.1} MB VRAM, {} total bars", + num_folds, wf_data.vram_bytes as f64 / 1_048_576.0, wf_data.total_bars, + ); + + // Store GPU walk-forward data and cudarc buffers for the experience collector + self.features_raw_cuda = Some(wf_data.features); + self.targets_raw_cuda = Some(wf_data.targets); + + let mut last_metrics = TrainingMetrics::new(); + + for fold_idx in 0..num_folds { + let fold = wf_data.folds.get(fold_idx).ok_or_else(|| { + anyhow::anyhow!("Fold {fold_idx} out of range") + })?; + + info!( + "=== Walk-Forward Fold {}/{} === train: {} bars, val: {} bars, test: {} bars", + fold_idx + 1, num_folds, fold.train_len(), fold.val_len(), fold.test_len(), + ); + + // Split training_data into fold's train and val slices (for CPU-side data) + let fold_train = &training_data[fold.train_start..fold.train_end]; + let fold_val: Vec<(FeatureVector, Vec)> = + training_data[fold.val_start..fold.val_end].to_vec(); + + // Store validation data for this fold + self.ofi_val_offset = fold.train_end; + self.val_data = fold_val; + + // Reset training state for new fold + self.gpu_data = None; // Force re-upload via DqnGpuData for the fold's range + self.best_sharpe = f64::NEG_INFINITY; + self.best_val_loss = f64::INFINITY; + self.loss_history.clear(); + self.q_value_history.clear(); + self.val_loss_history.clear(); + self.sharpe_history.clear(); + + // Run training loop on this fold's data + last_metrics = self + .train_with_data_full_loop(fold_train, &mut checkpoint_callback) + .await?; + + info!( + "Fold {}/{} complete: loss={:.6}, epochs={}", + fold_idx + 1, num_folds, + last_metrics.loss, + last_metrics.epochs_trained, + ); + } + + // Clean up GPU walk-forward buffers (features/targets already stored in self) + self.gpu_walk_forward = None; + + Ok(last_metrics) + } + + + + /// Calculate adaptive bounds with margin + /// + /// # Arguments + /// + /// * `stats` - Q-value statistics from Phase 1 + /// * `margin` - Safety margin as fraction (e.g., 0.3 = 30%) + /// + /// # Returns + /// + /// Tuple of (v_min, v_max) with safety margin applied + fn calculate_adaptive_bounds(stats: &QValueStats, margin: f64) -> (f64, f64) { + let range = stats.max - stats.min; + let v_min = stats.min - range * margin; + let v_max = stats.max + range * margin; + // Cap at ±10,000 to prevent explosion + (v_min.max(-10000.0), v_max.min(10000.0)) + } + + /// Reinitialize categorical distribution with new bounds + async fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), crate::MLError> { + let mut agent = self.agent.write().await; + match &mut *agent { + DQNAgentType::Standard(agent) => agent.reinit_categorical_distribution(v_min, v_max)?, + DQNAgentType::RegimeConditional(agent) => agent.reinit_categorical_distribution(v_min, v_max)?, + } + Ok(()) + } + + + + + + + + + + + + + + + + /// Calculate reward based on price movement + /// + /// # Arguments + /// * `current_close` - Current bar's close price + /// * `next_close` - Next bar's close price (target) + /// + /// # Returns + /// Normalized reward in [-1.0, 1.0] based on price change + fn calculate_reward(&self, current_close: f64, next_close: f64) -> f32 { + let price_change = next_close - current_close; + // Normalize by 10.0 for ES futures typical moves (±10 points) + // Clamp to [-1.0, 1.0] to prevent extreme rewards + (price_change / 10.0).clamp(-1.0, 1.0) as f32 + } + + /// Check if we can train (buffer has enough samples) + async fn can_train(&self) -> Result { + let agent = self.agent.read().await; + Ok(agent.can_train()) + } + + + + + /// Get current epsilon value + async fn get_epsilon(&self) -> Result { + let agent = self.agent.read().await; + Ok(agent.get_epsilon() as f64) + } + + /// Set epsilon value (used for deterministic evaluation) + async fn set_epsilon(&self, epsilon: f64) -> Result<()> { + let mut agent = self.agent.write().await; + agent.set_epsilon(epsilon); + Ok(()) + } + + /// Get best validation loss achieved during training + /// + /// Returns the lowest validation loss seen across all epochs. + /// Used by hyperopt adapter to optimize for generalization. + pub fn get_best_val_loss(&self) -> f64 { + self.best_val_loss + } + + /// Get epoch number where best validation loss was achieved + /// + /// Returns the 1-indexed epoch number with the best validation loss. + pub fn get_best_epoch(&self) -> usize { + self.best_epoch + } + + /// Get validation data for backtest integration + /// + /// Returns a reference to the validation dataset for hyperopt backtest evaluation. + /// Each entry contains a FeatureVector (42 market + 3 portfolio = 45 dims) and the corresponding target values. + /// Used by hyperopt adapter to run backtests on unseen data after training. + pub fn get_val_data(&self) -> &[(FeatureVector, Vec)] { + &self.val_data + } + + + + + /// Update portfolio tracker to reflect current position from backtest engine. + /// + /// Called between chunks so the next chunk's portfolio features + /// accurately reflect the current position (direction, value, exposure). + pub fn set_portfolio_for_backtest( + &mut self, + position_size: f32, + entry_price: f32, + current_price: f32, + ) { + self.portfolio_tracker = PortfolioTracker::new( + self.portfolio_tracker.initial_capital(), + self.portfolio_tracker.spread(), + 0.0, + ); + if position_size.abs() > f32::EPSILON { + self.portfolio_tracker + .set_position_direct(position_size, entry_price, current_price); + } + } + + /// Get the device used by this trainer + pub fn device(&self) -> &candle_core::Device { + &self.device + } + + /// Get access to the DQN agent + /// + /// Returns a reference to the Arc> for checkpoint saving. + /// Used by hyperopt adapter to save model weights after training. + pub fn get_agent(&self) -> &Arc> { + &self.agent + } + + /// Get reference to training hyperparameters + /// + /// Returns a reference to the DQN hyperparameters used for this trainer. + /// Used by tests to validate configuration. + pub fn hyperparams(&self) -> &DQNHyperparameters { + &self.hyperparams + } + + /// Get current learning rate from scheduler + /// + /// Returns the current learning rate after applying warmup and decay. + /// Used by tests and monitoring to track LR schedule. + pub fn get_current_lr(&self) -> f64 { + self.lr_scheduler.get_lr() + } + + /// Serialize model to bytes with architecture metadata embedded in safetensors header. + /// + /// For RegimeConditional agents, serializes ALL 3 heads (trending, ranging, + /// volatile) into a single safetensors file using prefixed tensor names + /// (`trending__`, `ranging__`, `volatile__`). This ensures walk-forward + /// checkpoint restore loads all heads, not just the trending head. + pub async fn serialize_model(&self) -> Result> { + let agent = self.agent.read().await; + + let tensors: std::collections::HashMap = match &*agent { + crate::trainers::dqn::DQNAgentType::RegimeConditional(regime) => { + let mut all_tensors = std::collections::HashMap::new(); + for (prefix, head_opt) in [ + ("trending__", regime.get_trending_head()), + ("ranging__", regime.get_ranging_head()), + ("volatile__", regime.get_volatile_head()), + ] { + let head = head_opt.ok_or_else(|| { + anyhow::anyhow!("Missing {} head for serialization", prefix) + })?; + let vars = head.get_q_network_vars(); + let vars_data = vars.data().lock().map_err(|_| { + anyhow::anyhow!("Failed to lock VarMap for {} head", prefix) + })?; + for (name, var) in vars_data.iter() { + all_tensors.insert( + format!("{}{}", prefix, name), + var.as_tensor().clone(), + ); + } + } + all_tensors + } + _ => { + let vars = agent.get_q_network_vars(); + let vars_data = vars.data().lock().map_err(|_| { + anyhow::anyhow!("Failed to lock VarMap for serialization") + })?; + vars_data + .iter() + .map(|(name, var)| (name.clone(), var.as_tensor().clone())) + .collect() + } + }; + + // Embed architecture metadata in safetensors header + let arch_metadata = Some(agent.checkpoint_metadata()); + let data = safetensors::serialize(&tensors, &arch_metadata) + .map_err(|e| anyhow::anyhow!("Failed to serialize safetensors: {}", e))?; + + Ok(data) + } + + /// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`). + /// + /// The trainer's `train_epoch` lazily uploads data on first call. + /// Use this to provide data that was uploaded in advance by a + /// `DoubleBufferedLoader`, skipping the per-fold upload latency. + pub fn set_gpu_data(&mut self, data: DqnGpuData) { + info!( + "DqnTrainer: injected pre-uploaded GPU data ({} bars, {:.1} MB)", + data.num_bars, + data.vram_bytes() as f64 / 1_048_576.0, + ); + self.gpu_data = Some(data); + } + + /// Drop cached GPU data, freeing VRAM for the next fold. + pub fn clear_gpu_data(&mut self) { + if self.gpu_data.is_some() { + info!("DqnTrainer: cleared GPU data (VRAM freed)"); + self.gpu_data = None; + } + } + + /// BUG #38 FIX: Clear replay buffer of contaminated experiences + pub async fn clear_replay_buffer(&mut self) -> Result<()> { + let mut agent = self.agent.write().await; + agent.clear_replay_buffer().map_err(|e| { + anyhow::anyhow!("Failed to clear replay buffer: {}", e) + })?; + let buffer_size = agent.get_replay_buffer_size().unwrap_or(0); + info!("Replay buffer cleared successfully. Current size: {}", buffer_size); + Ok(()) + } + + /// BUG #38 FIX: Reset target network to match current network + pub async fn reset_target_network(&mut self) -> Result<()> { + let mut agent = self.agent.write().await; + agent.reset_target_network().map_err(|e| { + anyhow::anyhow!("Failed to reset target network: {}", e) + })?; + info!("Target network reset successfully"); + Ok(()) + } + + + + + /// Get per-epoch training loss history (for smoke test verification) + pub fn loss_history(&self) -> &[f64] { + &self.loss_history + } + + /// Get per-epoch validation loss history + pub fn val_loss_history(&self) -> &[f64] { + &self.val_loss_history + } + + /// Get current epsilon from the DQN agent + pub async fn get_agent_epsilon(&self) -> f32 { + let agent_lock = self.agent.read().await; + agent_lock.get_epsilon() + } +} + +// --------------------------------------------------------------------------- +// GPU Q-value diagnostics (Task 6) +// --------------------------------------------------------------------------- + + + diff --git a/crates/ml/src/trainers/dqn/trainer/state.rs b/crates/ml/src/trainers/dqn/trainer/state.rs new file mode 100644 index 000000000..259d84395 --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/state.rs @@ -0,0 +1,179 @@ +//! DQN Trainer — State/feature vector conversion + +use anyhow::{Context, Result}; +use candle_core::Tensor; +use common::CommonError; +use num_traits::ToPrimitive; + +use super::DQNTrainer; +use crate::dqn::TradingState; +use crate::features::extraction::FeatureVector; + +impl DQNTrainer { + /// Convert feature vector to TradingState (42 market features → 45-dim state with portfolio) + /// + /// CRITICAL BUG FIX: Features 0-3 are LOG RETURNS (signed), not raw prices. + /// Using .abs() destroys directional information (bullish vs bearish). + /// We now use TradingState::from_normalized() to preserve sign information. + /// + /// Feature mapping: + /// - Features 0-3: OHLC log returns → price_features (signed, normalized) + /// - Features 4-224: All other features → technical_indicators (221 features including Wave D) + /// + /// # Arguments + /// + /// * `feature_vec` - 42-dimensional FeatureVector from extraction pipeline + /// * `close_price` - Current close price for portfolio feature calculation (optional) + /// + /// # Bug #4 Fix + /// + /// Added close_price parameter to enable portfolio feature population from PortfolioTracker. + pub(crate) fn feature_vector_to_state( + &self, + feature_vec: &FeatureVector, + close_price: Option, + ) -> Result { + self.feature_vector_to_state_with_ofi(feature_vec, close_price, None) + } + + pub(crate) fn feature_vector_to_state_with_ofi( + &self, + feature_vec: &FeatureVector, + close_price: Option, + ofi_index: Option, + ) -> Result { + // States are pre-normalized during data loading + let normalized_features: Vec = feature_vec.iter().map(|&v| v as f32).collect(); + + // Features 0-3 are LOG RETURNS - preserve sign information for price direction + let price_features: Vec = vec![ + normalized_features[0], // open log return (can be negative) + normalized_features[1], // high log return (can be negative) + normalized_features[2], // low log return (can be negative) + normalized_features[3], // close log return (can be negative) + ]; + + // 42-FEATURE ARCHITECTURE: Extract market features (indices 4-41) + assert_eq!( + normalized_features.len(), + 42, + "Expected 42 market features (got {})", + normalized_features.len() + ); + let market_features: Vec = normalized_features[4..42] + .iter() + .map(|&x| x as f32) + .collect(); + + // Legacy technical_indicators (empty for 42-feature architecture) + let technical_indicators = vec![]; + + // BUG #36 FIX: Use NORMALIZED portfolio features to prevent Q-value explosion + let portfolio_features = if let Some(price) = close_price { + let price_f32 = price.to_f32().unwrap_or(0.0); + self.portfolio_tracker + .get_portfolio_features(price_f32) + .to_vec() + } else { + vec![0.0, 0.0, 0.0] // Fallback if no price provided + }; + + // OFI regime features: 8 features from MBP-10 order book data. + // When OFI is enabled (mbp10_data_dir set), always return 8 features + // (zeros if data didn't load) to match state_dim=53. + let ofi_enabled = self.hyperparams.mbp10_data_dir.is_some(); + let regime_features: Vec = if let (Some(ofi), Some(idx)) = (&self.ofi_features, ofi_index) { + ofi.get(idx) + .map(|f| f.iter().map(|&v| v as f32).collect()) + .unwrap_or_else(|| vec![0.0; 8]) + } else if ofi_enabled { + vec![0.0; 8] + } else { + vec![] + }; + + // Use from_normalized() to preserve sign information + Ok(TradingState::from_normalized( + price_features, + technical_indicators, + market_features, + portfolio_features, + regime_features, + )) + } + + /// Convert feature vector to state tensor for action selection + /// + /// Public wrapper around internal state conversion for hyperopt backtest integration. + /// Converts a 42-dimensional feature vector to a 45-dimensional state tensor + /// suitable for DQN agent's select_action method. + /// + /// # Arguments + /// + /// * `feature_vec` - 42-dimensional market feature vector + /// * `close_price` - Current close price for portfolio feature calculation + /// + /// # Returns + /// + /// Result containing the 45-dimensional state tensor ready for model inference. + /// Portfolio features (last 3 dimensions) are populated via PortfolioTracker. + pub fn convert_to_state( + &self, + feature_vec: &FeatureVector, + close_price: f64, + ) -> Result { + let close = rust_decimal::Decimal::try_from(close_price) + .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; + + // Use internal conversion method (returns TradingState) + let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?; + + // Convert TradingState to flat vector, pad for tensor core alignment + let state_vec = trading_state.to_vector(); + let raw_dim = state_vec.len(); + let aligned = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); + let padded: Vec = if aligned > raw_dim { + let mut v = state_vec.to_vec(); + v.resize(aligned, 0.0); + v + } else { + state_vec.to_vec() + }; + + // Convert to Tensor using trainer's device (GPU or CPU) + Tensor::new(padded.as_slice(), &self.device) + .context("Failed to create state tensor from TradingState") + } + + /// Convert feature vector to flat state Vec (CPU only, no GPU tensor). + /// + /// Same as `convert_to_state` but returns the raw vector instead of a GPU tensor. + /// Used by chunked batch inference to avoid per-bar GPU allocations. + pub fn convert_to_state_vec( + &self, + feature_vec: &FeatureVector, + close_price: f64, + ) -> Result> { + let close = rust_decimal::Decimal::try_from(close_price) + .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; + let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?; + Ok(trading_state.to_vector()) + } + + /// Convert feature vector to flat state Vec with OFI features at the given index. + /// + /// Same as `convert_to_state_vec` but injects OFI features from the preloaded + /// array at `ofi_index`, preventing train/eval feature mismatch. + pub fn convert_to_state_vec_with_ofi( + &self, + feature_vec: &FeatureVector, + close_price: f64, + ofi_index: usize, + ) -> Result> { + let close = rust_decimal::Decimal::try_from(close_price) + .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?; + let trading_state = self.feature_vector_to_state_with_ofi(feature_vec, Some(close), Some(ofi_index))?; + Ok(trading_state.to_vector()) + } + +} diff --git a/crates/ml/src/trainers/dqn/trainer/tests.rs b/crates/ml/src/trainers/dqn/trainer/tests.rs new file mode 100644 index 000000000..363361a7e --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/tests.rs @@ -0,0 +1,609 @@ +use super::*; +use crate::hyperopt::ParameterSpace; +use std::sync::OnceLock; + +/// Shared CUDA device across all trainer tests. +/// +/// Root cause fix for CUBLAS_STATUS_NOT_INITIALIZED cascades: each +/// `Device::cuda_if_available(0)` creates a new cuBLAS handle. With +/// 400+ tests doing this in rapid succession (even with --test-threads=1), +/// the driver's internal handle pool is exhausted. Sharing one device +/// eliminates the churn entirely. +static SHARED_DEVICE: OnceLock = OnceLock::new(); + +fn shared_cuda_device() -> Device { + // Initialize tracing so kernel compilation/launch logs are visible. + static TRACING_INIT: std::sync::Once = std::sync::Once::new(); + TRACING_INIT.call_once(|| { + let filter = tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); + tracing_subscriber::fmt() + .with_env_filter(filter) + .with_test_writer() + .try_init() + .ok(); + }); + SHARED_DEVICE + .get_or_init(|| { + Device::cuda_if_available(0).unwrap_or(Device::Cpu) + }) + .clone() +} + +// Helper function to create test hyperparameters +// Uses conservative defaults suitable for testing +fn create_test_params() -> DQNHyperparameters { + let mut params = DQNHyperparameters::conservative(); + // Production default: branching DQN (3 heads: exposure, order, urgency). + // Always enabled — the warp-cooperative kernel on H100 requires it. + params.use_branching = true; + params.hidden_dim_base = Some(32); // Small for fast test iterations + params.buffer_size = 10_000; + params +} + +fn create_test_trainer() -> Result { + DQNTrainer::new_with_device(create_test_params(), shared_cuda_device()) +} + +fn create_test_trainer_with(params: DQNHyperparameters) -> Result { + DQNTrainer::new_with_device(params, shared_cuda_device()) +} + +/// Pad a TradingState's regime_features so that `state.dimension()` matches the +/// trainer's aligned state_dim (e.g. 45→48 on CUDA due to tensor core alignment). +fn pad_state_to_aligned(state: &mut TradingState, trainer: &DQNTrainer) { + let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores( + state.dimension(), + &trainer.device, + ); + let pad = aligned_dim.saturating_sub(state.dimension()); + if pad > 0 { + state.regime_features.extend(vec![0.0_f32; pad]); + } +} + +#[tokio::test] +async fn test_dqn_trainer_creation() { + let hyperparams = create_test_params(); + let trainer = create_test_trainer_with(hyperparams); + + assert!( + trainer.is_ok(), + "Failed to create DQN trainer: {:?}", + trainer.err() + ); +} + +#[tokio::test] +async fn test_batch_size_validation() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 500; + + // VRAM ceiling clamps if needed, never rejects + let trainer = create_test_trainer_with(hyperparams); + assert!( + trainer.is_ok(), + "Should clamp oversized batch, not reject: {:?}", + trainer.err() + ); +} + +#[tokio::test] +async fn test_feature_vector_to_state() { + let hyperparams = create_test_params(); + let trainer = create_test_trainer_with(hyperparams).unwrap(); + + // Create a synthetic 42-dim feature vector (42 market features) + let mut feature_vec = [0.0; 42]; + feature_vec[0] = 4000.0; // open + feature_vec[1] = 4010.0; // high + feature_vec[2] = 3990.0; // low + feature_vec[3] = 4005.0; // close + feature_vec[4] = 1000.0; // volume + // Fill remaining features with synthetic data + for i in 5..42 { + feature_vec[i] = (i as f64) * 0.1; + } + + let close_price = + rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); + let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)); + + assert!( + state.is_ok(), + "Failed to convert feature vector: {:?}", + state.err() + ); + + let state = state.unwrap(); + // State dimension is 45 (42 market + 3 portfolio) + // - Market features: 0-41 (42 features) + // - Portfolio features: 42-44 (3 features, populated by PortfolioTracker) + assert_eq!( + state.dimension(), + 45, + "State dimension should be 45 (42 market + 3 portfolio features)" + ); +} + +#[tokio::test] +async fn test_batched_action_selection() { + let hyperparams = create_test_params(); + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + // Create multiple synthetic states for batched action selection + let batch_size = 10; + let mut states = Vec::with_capacity(batch_size); + + for i in 0..batch_size { + let mut feature_vec = [0.0; 42]; // 42 market features + // Create varied states for testing + feature_vec[0] = 4000.0 + (i as f64 * 10.0); // open + feature_vec[1] = 4010.0 + (i as f64 * 10.0); // high + feature_vec[2] = 3990.0 + (i as f64 * 10.0); // low + feature_vec[3] = 4005.0 + (i as f64 * 10.0); // close + feature_vec[4] = 1000.0 + (i as f64 * 100.0); // volume + + // Fill remaining features + for j in 5..42 { + feature_vec[j] = (j as f64 + i as f64) * 0.1; + } + + let close_price = rust_decimal::Decimal::try_from(feature_vec[3]) + .unwrap_or(rust_decimal::Decimal::ZERO); + let mut state = trainer + .feature_vector_to_state(&feature_vec, Some(close_price)) + .unwrap(); + pad_state_to_aligned(&mut state, &trainer); + states.push(state); + } + + // Test batched action selection + let actions_result = trainer.select_actions_batch(&states).await; + + assert!( + actions_result.is_ok(), + "Batched action selection failed: {:?}", + actions_result.err() + ); + + let actions = actions_result.unwrap(); + assert_eq!( + actions.len(), + batch_size, + "Expected {} actions, got {}", + batch_size, + actions.len() + ); + + // Verify all actions have valid exposure indices (0-4) + for (i, action) in actions.iter().enumerate() { + let exp_idx = action.exposure as usize; + assert!( + exp_idx < 5, + "Action {} has invalid exposure index {}: {:?}", + i, + exp_idx, + action + ); + } +} + +#[tokio::test] +async fn test_batched_vs_sequential_action_selection_consistency() { + let hyperparams = create_test_params(); + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + // Create test states + let batch_size = 5; + let mut states = Vec::with_capacity(batch_size); + + for i in 0..batch_size { + let mut feature_vec = [0.0; 42]; // 42 market features + feature_vec[0] = 4000.0 + (i as f64 * 50.0); + feature_vec[1] = 4050.0 + (i as f64 * 50.0); + feature_vec[2] = 3950.0 + (i as f64 * 50.0); + feature_vec[3] = 4025.0 + (i as f64 * 50.0); + feature_vec[4] = 5000.0 + (i as f64 * 500.0); + + for j in 5..42 { + feature_vec[j] = (j as f64) * 0.5 + (i as f64); + } + + let close_price = rust_decimal::Decimal::try_from(feature_vec[3]) + .unwrap_or(rust_decimal::Decimal::ZERO); + let mut state = trainer + .feature_vector_to_state(&feature_vec, Some(close_price)) + .unwrap(); + pad_state_to_aligned(&mut state, &trainer); + states.push(state); + } + + // Get batched actions (GPU-optimized) + let batched_actions = trainer.select_actions_batch(&states).await.unwrap(); + + // Both should return valid actions + assert_eq!( + batched_actions.len(), + batch_size, + "Batched action count mismatch" + ); + + // Verify all actions have valid exposure indices (0-4) + for action in &batched_actions { + let exp_idx = action.exposure as usize; + assert!(exp_idx < 5, "Invalid exposure index {}: {:?}", exp_idx, action); + } +} + +#[tokio::test] +async fn test_empty_batch_handling() { + let hyperparams = create_test_params(); + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + let empty_states: Vec = Vec::new(); + let result = trainer.select_actions_batch(&empty_states).await; + + assert!(result.is_ok(), "Empty batch should be handled gracefully"); + assert_eq!( + result.unwrap().len(), + 0, + "Empty batch should return empty actions" + ); +} + +#[tokio::test] +async fn test_zero_batch_size_handling() { + // Test DQN rejects zero batch size + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 0; + + let result = create_test_trainer_with(hyperparams); + + // Should fail with descriptive error + assert!( + result.is_err(), + "DQN should reject zero batch size, but got: {:?}", + result + ); + + // Error message should mention batch size + let error_msg = result.unwrap_err().to_string(); + assert!( + error_msg.to_lowercase().contains("batch"), + "Error message should mention batch size, got: {}", + error_msg + ); +} + +// ===== Agent 23 Test #6: Batch Size Mismatch Validation Tests ===== + +/// Production-critical test: Verify trainer handles batch smaller than configured +#[tokio::test] +async fn test_batch_size_mismatch_smaller_than_configured() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 32; + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + // Create batch with 16 states (half of configured 32) + let mut feature_vec = [0.0; 42]; // 42 market features + for i in 0..4 { + feature_vec[i] = 4000.0 + (i as f64 * 10.0); + } + for i in 5..42 { + feature_vec[i] = (i as f64) * 0.1; + } + + let close_price = + rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); + let mut state = trainer + .feature_vector_to_state(&feature_vec, Some(close_price)) + .unwrap(); + pad_state_to_aligned(&mut state, &trainer); + let smaller_batch = vec![state.clone(); 16]; + + let result = trainer.select_actions_batch(&smaller_batch).await; + assert!( + result.is_ok(), + "DQN should handle smaller batches: {:?}", + result.err() + ); + assert_eq!( + result.unwrap().len(), + 16, + "Should return action for each state" + ); +} + +/// Production-critical test: Verify trainer handles batch larger than configured +#[tokio::test] +async fn test_batch_size_mismatch_larger_than_configured() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 16; + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + // Create batch with 64 states (4x configured 16) + let mut feature_vec = [0.0; 42]; // 42 market features + for i in 0..4 { + feature_vec[i] = 4000.0 + (i as f64 * 10.0); + } + for i in 5..42 { + feature_vec[i] = (i as f64) * 0.1; + } + + let close_price = + rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); + let mut state = trainer + .feature_vector_to_state(&feature_vec, Some(close_price)) + .unwrap(); + pad_state_to_aligned(&mut state, &trainer); + let larger_batch = vec![state.clone(); 64]; + + let result = trainer.select_actions_batch(&larger_batch).await; + assert!( + result.is_ok(), + "DQN should handle larger batches: {:?}", + result.err() + ); + assert_eq!( + result.unwrap().len(), + 64, + "Should return action for each state" + ); +} + +/// Production-critical test: Verify empty batch handling +#[tokio::test] +async fn test_empty_batch_returns_empty_actions() { + let mut trainer = create_test_trainer().unwrap(); + let empty_batch: Vec = vec![]; + + let result = trainer.select_actions_batch(&empty_batch).await; + assert!(result.is_ok(), "Should handle empty batch gracefully"); + assert_eq!( + result.unwrap().len(), + 0, + "Empty batch should return empty actions" + ); +} + +/// Production-critical test: Verify single-sample batch handling +#[tokio::test] +async fn test_single_sample_batch() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 32; + let mut trainer = create_test_trainer_with(hyperparams).unwrap(); + + let mut feature_vec = [0.0; 42]; // 42 market features + for i in 0..4 { + feature_vec[i] = 4000.0; + } + for i in 5..42 { + feature_vec[i] = (i as f64) * 0.1; + } + + let close_price = + rust_decimal::Decimal::try_from(feature_vec[3]).unwrap_or(rust_decimal::Decimal::ZERO); + let mut state = trainer + .feature_vector_to_state(&feature_vec, Some(close_price)) + .unwrap(); + pad_state_to_aligned(&mut state, &trainer); + let single_batch = vec![state]; + + let result = trainer.select_actions_batch(&single_batch).await; + assert!( + result.is_ok(), + "Should handle single-sample batch: {:?}", + result.err() + ); + assert_eq!(result.unwrap().len(), 1, "Should return exactly one action"); +} + +/// Large batch sizes are accepted (VRAM ceiling is the only cap) +#[test] +fn test_large_batch_size_accepted() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 2048; + + let result = create_test_trainer_with(hyperparams); + assert!( + result.is_ok(), + "Should accept large batch sizes within VRAM ceiling: {:?}", + result.err() + ); +} + +/// Production-critical test: Non-power-of-2 batch sizes +#[tokio::test] +async fn test_non_power_of_two_batch_size() { + let mut hyperparams = create_test_params(); + hyperparams.batch_size = 13; // Not a power of 2 + + let result = create_test_trainer_with(hyperparams); + assert!( + result.is_ok(), + "Should accept non-power-of-2 batch sizes: {:?}", + result.err() + ); +} + +/// Production-critical test: Train with empty dataset doesn't crash +#[tokio::test] +async fn test_train_with_empty_data_completes_gracefully() { + let mut params = create_test_params(); + params.epochs = 5; // Short run — just checking it doesn't panic + params.early_stopping_enabled = false; + params.gradient_collapse_patience = 1000; + params.buffer_size = 1000; + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let mut trainer = DQNTrainer::new_with_device(params, device).unwrap(); + let empty_data: Vec<(FeatureVector, Vec)> = vec![]; + let checkpoint_callback = |_, _, _| Ok(String::new()); + + let result = trainer + .train_with_data_full_loop(&empty_data, checkpoint_callback) + .await; + + assert!( + result.is_err(), + "Training with empty data should return an error (no CPU fallback)" + ); +} + +/// Test reward function calculates actual price changes correctly +#[test] +fn test_reward_function_price_changes() { + let trainer = create_test_trainer().unwrap(); + + // Test upward price move (+14.25 points, should clamp to +1.0) + let reward_up = trainer.calculate_reward(5900.0, 5914.25); + assert!( + (reward_up - 1.0).abs() < 1e-6, + "Upward move should return +1.0 (clamped), got: {}", + reward_up + ); + + // Test downward price move (-14.25 points, should clamp to -1.0) + let reward_down = trainer.calculate_reward(5914.25, 5900.0); + assert!( + (reward_down - (-1.0)).abs() < 1e-6, + "Downward move should return -1.0 (clamped), got: {}", + reward_down + ); + + // Test flat market (0 points, should return 0.0) + let reward_flat = trainer.calculate_reward(5900.0, 5900.0); + assert!( + reward_flat.abs() < 1e-6, + "Flat market should return 0.0, got: {}", + reward_flat + ); + + // Test small upward move (+5 points, should return +0.5) + let reward_small_up = trainer.calculate_reward(5900.0, 5905.0); + assert!( + (reward_small_up - 0.5).abs() < 1e-6, + "Small upward move (+5) should return +0.5, got: {}", + reward_small_up + ); + + // Test small downward move (-5 points, should return -0.5) + let reward_small_down = trainer.calculate_reward(5905.0, 5900.0); + assert!( + (reward_small_down - (-0.5)).abs() < 1e-6, + "Small downward move (-5) should return -0.5, got: {}", + reward_small_down + ); + + // Test unclamped move (+3 points, should return +0.3) + let reward_unclamped = trainer.calculate_reward(5900.0, 5903.0); + assert!( + (reward_unclamped - 0.3).abs() < 1e-6, + "Move of +3 points should return +0.3, got: {}", + reward_unclamped + ); +} + +#[test] +fn test_dynamic_batch_size_l4() { + // L4 has 24GB VRAM — HardwareBudget should allow batch_size >> 230 + let budget = crate::hyperopt::HardwareBudget { + gpu_memory_mb: 24_000, + gpu_name: "NVIDIA L4".to_string(), + }; + let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); + assert!(batch.unwrap_or(0.0) > 230.0, "L4 should support DQN batch > 230, got {:?}", batch); +} + +#[test] +fn test_dynamic_batch_size_h100() { + // H100 has 80GB VRAM — should hit the 8192 ceiling + let budget = crate::hyperopt::HardwareBudget { + gpu_memory_mb: 81_920, + gpu_name: "NVIDIA H100".to_string(), + }; + let batch = budget.max_batch_size(50.0, 0.0005, 64.0, 8192.0); + assert!((batch.unwrap_or(0.0) - 8192.0).abs() < 1.0, "H100 should hit 8192 ceiling, got {:?}", batch); +} + +// ── C2 Overhaul Smoke Tests ───────────────────────────────────────── + +/// Verify DQN action space is 5 exposure levels (not 45 factored actions). +#[test] +fn test_c2_dqn_default_num_actions_is_5() { + let config = crate::dqn::DQNConfig::default(); + assert_eq!(config.num_actions, 5, "DQN default must be 5 exposure-level actions"); +} + +/// Verify 5 exposure indices produce 5 distinct exposure levels. +#[test] +fn test_c2_five_actions_produce_distinct_exposures() { + use crate::dqn::action_space::ExposureLevel; + use crate::dqn::order_router::OrderRouter; + + let actions: Vec<_> = (0..5) + .filter_map(|idx| ExposureLevel::from_index(idx).ok()) + .map(|e| OrderRouter::route_default(e)) + .collect(); + + assert_eq!(actions.len(), 5); + + let unique: std::collections::HashSet<_> = actions.iter().map(|a| a.exposure).collect(); + assert_eq!(unique.len(), 5, "All 5 exposure levels must be distinct"); +} + +/// Verify hyperopt search space is 29D (C4: sharpe_weight, L2: branch_hidden_dim). +#[test] +fn test_c3_search_space_is_27d() { + let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds(); + assert_eq!(bounds.len(), 30, "Search space must be 30D (C6: gradient_accumulation_steps added)"); + + let names = crate::hyperopt::adapters::dqn::DQNParams::param_names(); + assert_eq!(names.len(), 30); + assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)"); + assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)"); + assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)"); + assert!(!names.contains(&"curiosity_weight"), "curiosity_weight must not be in search space"); + assert!(!names.contains(&"noisy_epsilon_floor"), "noisy_epsilon_floor must not be in search space"); +} + +/// Verify noisy_epsilon_floor is fixed to 0.10 (prevents action collapse). +#[test] +fn test_noisy_epsilon_floor_fixed() { + let params = crate::hyperopt::adapters::dqn::DQNParams::default(); + assert!( + (params.noisy_epsilon_floor - 0.10).abs() < 1e-6, + "noisy_epsilon_floor must default to 0.10 (prevents action collapse)" + ); +} + +/// Verify exploration params are fixed after C2 cleanup. +#[test] +fn test_c2_exploration_params_fixed() { + let params = crate::hyperopt::adapters::dqn::DQNParams::default(); + assert!( + params.curiosity_weight.abs() < f64::EPSILON, + "curiosity_weight must be fixed at 0.0" + ); + assert!( + (params.noisy_epsilon_floor - 0.10).abs() < 1e-6, + "noisy_epsilon_floor must be fixed at 0.10" + ); + assert!( + params.count_bonus_coefficient.abs() < f64::EPSILON, + "count_bonus_coefficient must be fixed at 0.0" + ); + + // Roundtrip through from_continuous should preserve fixed values + let continuous = params.to_continuous(); + let recovered = crate::hyperopt::adapters::dqn::DQNParams::from_continuous(&continuous).unwrap(); + assert!( + recovered.curiosity_weight.abs() < f64::EPSILON, + "curiosity_weight must remain 0.0 after roundtrip" + ); + assert!( + recovered.count_bonus_coefficient.abs() < f64::EPSILON, + "count_bonus_coefficient must remain 0.0 after roundtrip" + ); +} diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs new file mode 100644 index 000000000..a60020979 --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -0,0 +1,631 @@ +//! DQN training step methods — single-batch and gradient-accumulation paths. + +use anyhow::Result; +use candle_core::{IndexOp, Tensor}; +use tracing::{debug, info, warn}; + +use crate::dqn::mixed_precision::training_dtype; +use super::DQNTrainer; + +impl DQNTrainer { + /// Perform one training step using real DQN algorithm + /// + /// This method implements the core Deep Q-Learning algorithm: + /// 1. Sample batch from experience replay buffer + /// 2. Compute current Q-values: Q(s, a) + /// 3. Compute target Q-values: r + γ * max_a' Q_target(s', a') + /// 4. Calculate TD-error and MSE loss + /// 5. Backpropagate gradients and update Q-network + /// 6. Periodically update target network + /// + /// WAVE 26 P2.2: Now supports gradient accumulation for larger effective batch sizes + /// + /// Returns: (loss, avg_q_value, grad_norm) + async fn train_step(&mut self) -> Result<(f64, f64, f64)> { + let accumulation_steps = self.hyperparams.gradient_accumulation_steps; + + // OOM recovery loop: retry up to 3 times with halved batch size + const MAX_OOM_RETRIES: usize = 3; + + for retry in 0..=MAX_OOM_RETRIES { + let result = if accumulation_steps > 1 { + self.train_step_with_accumulation().await + } else { + self.train_step_single_batch().await + }; + + match result { + Ok(metrics) => return Ok(metrics), + Err(e) => { + // Check if this is an OOM error by inspecting the error chain + let err_str = format!("{:?}", e).to_lowercase(); + let is_oom = err_str.contains("out of memory") + || err_str.contains("oom") + || err_str.contains("cuda error 2") + || err_str.contains("cudamalloc") + || err_str.contains("failed to allocate"); + + if is_oom && retry < MAX_OOM_RETRIES { + let old_batch = self.current_batch_size; + self.current_batch_size = (old_batch / 2).max(1); + warn!( + "OOM detected (retry {}/{}): reducing batch size {} -> {}", + retry + 1, + MAX_OOM_RETRIES, + old_batch, + self.current_batch_size + ); + // Continue to next retry + } else { + return Err(e); + } + } + } + } + + Err(anyhow::anyhow!( + "Training failed after {} OOM retries", + MAX_OOM_RETRIES + )) + } + + /// Standard single-batch training step (no gradient accumulation) + /// + /// Pre-samples from PER buffer using READ lock before acquiring WRITE lock + /// for GPU training. This preserves PER IS-weights and indices for proper + /// importance sampling correction and priority updates. + async fn train_step_single_batch(&mut self) -> Result<(f64, f64, f64)> { + // Pre-sample batch OUTSIDE the write lock using read-only access to the buffer. + // PER IS-weights and indices are preserved for correct importance sampling. + // The write lock is only held during GPU forward/backward + optimizer step. + let explicit_batch = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + let sample_size = self.current_batch_size; + buffer.can_sample(sample_size).then(|| { + buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Failed to pre-sample batch: {}", e)) + }).transpose()? + }; // READ lock released here + + let mut agent = self.agent.write().await; + + // train_step returns GpuTrainResult with GPU-resident scalar tensors. + #[allow(unused_variables)] + let gpu_result = agent + .train_step(explicit_batch) + .map_err(|e| anyhow::anyhow!("Training step failed: {}", e))?; + + // GPU training guard: on-device NaN/loss-clip/grad-collapse checks. + // Zero cudaStreamSynchronize — kernel writes halt flags to pinned host memory. + let (loss_clipped, grad_norm) = { + // Lazy-init training guard on first call + if self.training_guard.is_none() && self.device.is_cuda() { + match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) { + Ok(guard) => { + info!("GPU training guard initialized"); + self.training_guard = Some(guard); + } + Err(e) => { + return Err(anyhow::anyhow!("GPU training guard init FAILED (no CPU fallback): {e}")); + } + } + } + + if let Some(ref mut guard) = self.training_guard { + let grad_collapse_threshold = + self.hyperparams.learning_rate as f32 + * self.hyperparams.gradient_collapse_multiplier as f32; + // Use original buffer_size (before AutoReplaySizer) for warmup guard + let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; + let past_warmup = self.gradient_logging_step as u64 > warmup_steps; + + let result = guard + .check_and_accumulate( + &gpu_result.loss_gpu, + &gpu_result.grad_norm_gpu, + 1e6_f32, // loss clip threshold + grad_collapse_threshold, + !past_warmup, + ) + .map_err(|e| anyhow::anyhow!("GPU guard check: {e}"))?; + + // Handle halt conditions + if result.halt_nan { + return Err(anyhow::anyhow!( + "NaN/Inf detected in loss ({}) or grad_norm ({})", + result.raw_loss, + result.raw_grad_norm + )); + } + if result.halt_loss_clip { + warn!( + "Loss clipped from {:.2e} to 1.0e6 (TD error explosion, epoch {})", + result.raw_loss, + self.loss_history.len() + 1 + ); + } + + // GPU guard path: collapse check only (no detect_dead_neurons GPU->CPU sync). + // Dead neuron detection runs at epoch boundary via log_diagnostics(). + agent.check_gradient_collapse(result.raw_grad_norm).map_err(|e| { + tracing::info!("Early stopping triggered (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + (result.clipped_loss as f64, result.raw_grad_norm as f64) + } else { + return Err(anyhow::anyhow!( + "GPU training guard not initialized — CUDA device required for DQN training" + )); + } + }; + // Unreachable in non-cuda mode (return above), but Rust still name-checks. + + // Q-value estimation: periodic (every 50 steps) + self.q_estimation_counter += 1; + if self.q_estimation_counter % 50 == 1 { + // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) + let mut gpu_q_done = false; + { + if let Some(ref mut guard) = self.training_guard { + let buffer = agent.memory(); + if buffer.len() > 0 { + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; + + let state_dim = agent.get_state_dim(); + let mut batch_tensor_opt: Option = None; + + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = Some( + gpu.states + .to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, + ); + } + if batch_tensor_opt.is_none() { + let mut state_data = + Vec::with_capacity(sample_size * state_dim); + for exp in &batch_sample.experiences { + state_data.extend_from_slice(&exp.state); + } + if !state_data.is_empty() { + let tensor = Tensor::from_vec( state_data, + (sample_size, state_dim), + &self.device, + ) + .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; + batch_tensor_opt = Some(tensor); + } + } + + if let Some(ref batch_tensor) = batch_tensor_opt { + // Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync + agent.set_training_forward_active(true); + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; + agent.set_training_forward_active(false); + let num_actions = + batch_q_values.dims().get(1).copied().unwrap_or(5); + + // Divergence check on first sample + let first_q = batch_q_values + .i(0) + .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; + let div_result = guard + .qvalue_divergence(&first_q, num_actions, 10000.0) + .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; + agent + .log_q_values_from_stats( + div_result.q_min, + div_result.q_max, + div_result.q_mean, + div_result.q_variance, + num_actions, + ) + .map_err(|e| { + tracing::info!( + "Early stopping (Q-value divergence): {}", + e + ); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Batch average via GPU reduction (one-step delay due to double-buffering) + let stats = guard + .qvalue_stats(&batch_q_values, sample_size, num_actions) + .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; + self.cached_avg_q = stats.q_mean as f64; + + // Accumulate Q-value mean on GPU via Welford running mean (zero sync) + let avg_q_tensor = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; + guard + .accumulate_q_value(&avg_q_tensor) + .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; + + gpu_q_done = true; + } + } + } + } + // CUDA: GPU Q-value accumulation is mandatory — no CPU fallback. + if !gpu_q_done { + return Err(anyhow::anyhow!( + "GPU Q-value accumulation FAILED (no CPU fallback). \ + Check GpuTrainingGuard initialization." + )); + } + } + let avg_q_value = self.cached_avg_q; + + debug!("Gradient norm after clip (actual): {:.4}", grad_norm); + + self.gradient_logging_step += 1; + if self.gradient_logging_step % 10 == 0 { + debug!( + "Step {}: grad={:.4}, loss={:.4}", + self.gradient_logging_step, grad_norm, loss_clipped + ); + } + + Ok((loss_clipped, avg_q_value, grad_norm)) + } + + /// Training step with true gradient accumulation across N mini-batches. + /// + /// Unlike the previous implementation which ran N independent optimizer + /// steps, this version computes gradients for each mini-batch, accumulates + /// them, averages, and then applies a **single** optimizer step. This + /// simulates training with an effective batch size of + /// `accumulation_steps * batch_size` while keeping memory usage at + /// `batch_size`. + /// + /// Returns: (avg_loss, avg_q_value, final_grad_norm) + async fn train_step_with_accumulation(&mut self) -> Result<(f64, f64, f64)> { + let accumulation_steps = self.hyperparams.gradient_accumulation_steps; + + debug!( + "Starting true gradient accumulation with {} steps (effective batch: {})", + accumulation_steps, + self.current_batch_size * accumulation_steps + ); + + // Pre-sample ALL mini-batches using READ lock (no GPU contention). + let pre_sampled: Vec> = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + let sample_size = self.current_batch_size; + let mut batches = Vec::with_capacity(accumulation_steps); + for step_idx in 0..accumulation_steps { + batches.push( + buffer.can_sample(sample_size).then(|| { + buffer.sample(sample_size).map_err(|e| { + anyhow::anyhow!( + "Failed to pre-sample batch (accum step {}): {}", + step_idx, + e + ) + }) + }).transpose()?, + ); + } + batches + }; // READ lock released + + let mut agent = self.agent.write().await; + + // === Phase 1: Accumulate gradients across N mini-batches === + let mut accumulated_grads: Option = None; + // Used by non-CUDA fallback and CUDA empty-tensor fallback paths. + #[allow(unused_mut, unused_assignments, unused_variables)] + let mut total_loss = 0.0_f64; + let mut all_td_errors = Vec::new(); + let mut all_indices = Vec::new(); + #[allow(unused_mut, unused_assignments, unused_variables)] + let mut final_grad_norm = 0.0_f32; + let mut gpu_td_errors: Vec = Vec::new(); + let mut gpu_indices: Vec = Vec::new(); + let mut gpu_loss_tensors: Vec = Vec::new(); + let mut gpu_grad_tensors: Vec = Vec::new(); + + for (step, batch) in pre_sampled.into_iter().enumerate() { + // Compute forward pass + backward WITHOUT optimizer step + let result = agent + .compute_gradients(batch) + .map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?; + + // Get vars for accumulation. Var is an Arc wrapper so cloning is cheap. + let vars: Vec = agent + .optimizer_vars() + .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?; + + crate::gradient_accumulation::accumulate_grads( + &mut accumulated_grads, + result.grads, + &vars, + ) + .map_err(|e| anyhow::anyhow!("Gradient accumulation step {} failed: {}", step, e))?; + + all_td_errors.extend(result.td_errors); + all_indices.extend(result.indices); + // GPU guard: check + accumulate loss/grad for this sub-step (borrows + // tensors before the move into gpu_*_tensors below). + { + if let (Some(ref loss_gpu), Some(ref gn_gpu)) = + (&result.loss_tensor_gpu, &result.grad_norm_gpu) + { + if let Some(ref mut guard) = self.training_guard { + let grad_collapse_threshold = + self.hyperparams.learning_rate as f32 + * self.hyperparams.gradient_collapse_multiplier as f32; + // Use original buffer_size (before AutoReplaySizer) for warmup guard + let warmup_steps = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; + let past_warmup = self.gradient_logging_step as u64 > warmup_steps; + + let guard_result = guard.check_and_accumulate( + loss_gpu, + gn_gpu, + 1e6_f32, + grad_collapse_threshold, + !past_warmup, + ).map_err(|e| anyhow::anyhow!("GPU guard sub-step {}: {e}", step))?; + + if guard_result.halt_nan { + return Err(anyhow::anyhow!( + "NaN/Inf at accumulation sub-step {}: loss={}, grad={}", + step, guard_result.raw_loss, guard_result.raw_grad_norm + )); + } + if guard_result.halt_grad_collapse { + agent.check_gradient_collapse(guard_result.raw_grad_norm).map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + } + } + } + { + if let Some(td_gpu) = result.td_errors_gpu { + gpu_td_errors.push(td_gpu); + } + if let Some(idx_gpu) = result.indices_gpu { + gpu_indices.push(idx_gpu); + } + if let Some(loss_gpu) = result.loss_tensor_gpu { + gpu_loss_tensors.push(loss_gpu); + } + if let Some(gn_gpu) = result.grad_norm_gpu { + gpu_grad_tensors.push(gn_gpu); + } + } + // CPU sentinel fallback (non-CUDA only) + } + + // === Phase 2: Average and apply gradients (single optimizer step) === + if let Some(ref mut grads) = accumulated_grads { + let vars: Vec = agent + .optimizer_vars() + .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?; + + crate::gradient_accumulation::scale_grads( + grads, + &vars, + 1.0 / accumulation_steps as f64, + ) + .map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?; + + let guard_active = self.training_guard.is_some(); + crate::gradient_accumulation::check_gradients_finite_guarded( + grads, + &vars, + guard_active, + ).map_err(|e| anyhow::anyhow!("Training halted: {}", e))?; + + agent + .apply_accumulated_gradients(grads) + .map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?; + } + + // === Phase 3: Bookkeeping === + { + // GPU PER path: concatenate GPU tensors and update in one shot + if !gpu_td_errors.is_empty() && !gpu_indices.is_empty() { + let td_cat = candle_core::Tensor::cat(&gpu_td_errors, 0) + .map_err(|e| anyhow::anyhow!("GPU TD error concat failed: {}", e))?; + let idx_cat = candle_core::Tensor::cat(&gpu_indices, 0) + .map_err(|e| anyhow::anyhow!("GPU index concat failed: {}", e))?; + agent + .update_priorities_gpu(&idx_cat, &td_cat) + .map_err(|e| anyhow::anyhow!("GPU PER priority update failed: {}", e))?; + } else if !all_indices.is_empty() { + agent + .update_priorities(&all_indices, &all_td_errors) + .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?; + } else { + // No priority updates needed (uniform buffer or empty batch) + } + } + agent.step_replay_buffer(); + + // Single readback at accumulation boundary — prefer GPU guard accumulators + // (zero extra sync), fall back to cat+mean+to_scalar if guard absent. + let (avg_loss, final_grad_norm_f64) = { + if let Some(ref mut guard) = self.training_guard { + let (avg_l, avg_gn) = guard.read_accumulators() + .map_err(|e| anyhow::anyhow!("GPU guard read_accumulators: {e}"))?; + guard.reset_accumulators() + .map_err(|e| anyhow::anyhow!("GPU guard reset: {e}"))?; + (avg_l, avg_gn) + } else { + return Err(anyhow::anyhow!( + "GPU training guard not initialized — CUDA device required" + )); + } + }; + + // Gradient collapse detection (early stopping) -- no dead neuron GPU->CPU sync. + // Dead neuron detection runs at epoch boundary via log_diagnostics(). + agent + .check_gradient_collapse(final_grad_norm_f64 as f32) + .map_err(|e| { + tracing::info!("Early stopping triggered (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Clip averaged loss + let loss_clipped = if avg_loss > 1e6 { + warn!("Averaged loss clipped from {:.2e} to 1.0e6", avg_loss); + 1e6 + } else { + avg_loss + }; + + // Q-value estimation: periodic (every 50 steps) + self.q_estimation_counter += 1; + if self.q_estimation_counter % 50 == 1 { + // GPU path: use qvalue_stats / qvalue_divergence kernels (zero to_scalar readback) + let mut gpu_q_done = false; + { + if let Some(ref mut guard) = self.training_guard { + let buffer = agent.memory(); + if buffer.len() > 0 { + let sample_size = buffer.len().min(10); + let batch_sample = buffer + .sample(sample_size) + .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; + + let state_dim = agent.get_state_dim(); + let mut batch_tensor_opt: Option = None; + + if let Some(ref gpu) = batch_sample.gpu_batch { + batch_tensor_opt = Some( + gpu.states + .to_dtype(training_dtype(agent.device())) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?, + ); + } + if batch_tensor_opt.is_none() { + let mut state_data = + Vec::with_capacity(sample_size * state_dim); + for exp in &batch_sample.experiences { + state_data.extend_from_slice(&exp.state); + } + if !state_data.is_empty() { + let tensor = Tensor::from_vec( state_data, + (sample_size, state_dim), + &self.device, + ) + .map_err(|e| anyhow::anyhow!("Q-est tensor: {e}"))? + .to_dtype(training_dtype(&self.device)) + .map_err(|e| anyhow::anyhow!("Q-est dtype: {e}"))?; + batch_tensor_opt = Some(tensor); + } + } + + if let Some(ref batch_tensor) = batch_tensor_opt { + // Suppress forward() monitoring to avoid to_vec2 GPU→CPU sync + agent.set_training_forward_active(true); + let batch_q_values = agent + .forward(batch_tensor) + .map_err(|e| anyhow::anyhow!("Q-est forward: {e}"))?; + agent.set_training_forward_active(false); + let num_actions = + batch_q_values.dims().get(1).copied().unwrap_or(5); + + // Divergence check on first sample + let first_q = batch_q_values + .i(0) + .map_err(|e| anyhow::anyhow!("Q-est index: {e}"))?; + let div_result = guard + .qvalue_divergence(&first_q, num_actions, 10000.0) + .map_err(|e| anyhow::anyhow!("GPU Q-div: {e}"))?; + agent + .log_q_values_from_stats( + div_result.q_min, + div_result.q_max, + div_result.q_mean, + div_result.q_variance, + num_actions, + ) + .map_err(|e| { + tracing::info!( + "Early stopping (Q-value divergence): {}", + e + ); + anyhow::anyhow!("Early stopping: {}", e) + })?; + + // Batch average via GPU reduction (one-step delay due to double-buffering) + let stats = guard + .qvalue_stats(&batch_q_values, sample_size, num_actions) + .map_err(|e| anyhow::anyhow!("GPU Q-stats: {e}"))?; + self.cached_avg_q = stats.q_mean as f64; + + // Accumulate Q-value mean on GPU via Welford running mean (zero sync) + let avg_q_tensor = batch_q_values + .max(1) + .map_err(|e| anyhow::anyhow!("GPU Q-acc max: {e}"))? + .mean_all() + .map_err(|e| anyhow::anyhow!("GPU Q-acc mean: {e}"))?; + guard + .accumulate_q_value(&avg_q_tensor) + .map_err(|e| anyhow::anyhow!("GPU Q-acc: {e}"))?; + + gpu_q_done = true; + } + } + } + } + // CUDA: GPU Q-value accumulation is mandatory — no CPU fallback. + if !gpu_q_done { + return Err(anyhow::anyhow!( + "GPU Q-value accumulation FAILED (no CPU fallback). \ + Check GpuTrainingGuard initialization." + )); + } + } + let avg_q_value = self.cached_avg_q; + + Ok((loss_clipped, avg_q_value, final_grad_norm_f64)) + } + + /// Lazy-init fused CUDA training context (Standard DQN only). + /// Recreate if batch_size changed (OOM recovery). + pub(crate) async fn ensure_fused_ctx(&mut self) { + let needs_init = match &self.fused_ctx { + None => self.device.is_cuda(), + Some(ctx) => ctx.batch_size() != self.current_batch_size, + }; + if !needs_init { + return; + } + if self.fused_ctx.is_some() { + info!("Fused CUDA context: batch_size changed, recreating"); + self.fused_ctx = None; + } + let agent = self.agent.read().await; + match super::super::fused_training::FusedTrainingCtx::new( + &self.device, &*agent, &self.hyperparams, self.current_batch_size, + ) { + Ok(ctx) => { + info!("Fused CUDA training initialized (batch_size={})", self.current_batch_size); + self.fused_ctx = Some(ctx); + } + Err(e) => { + warn!("Fused CUDA training init failed, using Candle path: {e}"); + } + } + } +} diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs new file mode 100644 index 000000000..e9e9cb2f3 --- /dev/null +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -0,0 +1,1847 @@ +//! DQN main training loop — `train_with_data_full_loop`. +//! +//! The main loop delegates to focused helper methods on `DQNTrainer`: +//! - `log_training_config`: one-time Rainbow/target-update logging +//! - `init_gpu_data`: upload training data to GPU (Phase 1) +//! - `init_gpu_raw_buffers`: upload raw cudarc targets/features + portfolio sim (Phase 1b) +//! - `init_gpu_experience_collector`: build the zero-roundtrip CUDA collector (Phase 1c) +//! - `collect_gpu_experiences`: run the GPU experience collection kernel (Phase 3) +//! - `run_training_steps`: batched training from replay buffer with guard kernel +//! - `process_epoch_boundary`: single readback + safety checks at epoch end +//! - `sync_gpu_weights`: push updated weights to GPU collector +//! - `refresh_stale_per_priorities`: M2 PER staleness refresh +//! - `log_epoch_metrics_and_financials`: logging, Prometheus, QuestDB, VaR/CVaR +//! - `handle_epoch_checkpoints_and_early_stopping`: best-checkpoint + early stopping + +use anyhow::{Context, Result}; +use candle_core::Tensor; +use common::metrics::{questdb_sink, training_metrics}; +use tracing::{debug, info, warn}; + +use crate::cuda_pipeline::DqnGpuData; +use crate::dqn::logging::{log_epoch_start, log_epoch_end, log_training_progress}; +use crate::dqn::target_update::convergence_half_life; +use crate::evaluation::metrics::calculate_var_cvar; +use crate::trainers::TargetUpdateMode; +use crate::TrainingMetrics; +use crate::features::extraction::FeatureVector; +use super::super::config::DQNAgentType; +use super::super::financials::compute_epoch_financials; +use super::super::monitoring::TrainingMonitor; +use super::{DQNTrainer, EPISODE_LENGTH}; + +/// Metrics returned from `process_epoch_boundary` for downstream logging. +pub(crate) struct EpochBoundaryMetrics { + pub avg_loss: f32, + pub avg_grad: f32, + pub avg_q: f64, +} + +/// Metrics returned from `log_epoch_metrics_and_financials` for checkpointing. +pub(crate) struct EpochLogOutput { + pub avg_loss: f64, + pub avg_q_value: f64, + pub avg_grad_norm: f64, + pub epoch_sharpe: f64, + pub val_loss: f64, + pub q_min: f64, + pub q_max: f64, + pub q_mean: f64, +} + +impl DQNTrainer { + // ═══════════════════════════════════════════════════════════════════════ + // Main training loop — orchestrates helpers + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn train_with_data_full_loop( + &mut self, + training_data: &[(FeatureVector, Vec)], + mut checkpoint_callback: F, + ) -> Result + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + let start_time = std::time::Instant::now(); + let mut total_loss = 0.0; + let mut total_q_value = 0.0; + let mut total_gradient_norm = 0.0; + let mut total_reward = 0.0; + let mut total_action_counts = [0_usize; 5]; + let mut total_factored_action_counts = [0_usize; 45]; + + self.log_training_config().await; + + // Training loop + for epoch in 0..self.hyperparams.epochs { + self.reset_epoch_state(epoch); + + log_epoch_start(epoch + 1, self.hyperparams.epochs, self.hyperparams.learning_rate); + training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64); + + let mut monitor = TrainingMonitor::new(epoch + 1); + + // BUG #15 FIX: Portfolio compounds across epochs (see detailed comment in reset_epoch_state) + self.portfolio_tracker.reset_drawdown_tracking(); + + let epoch_start = std::time::Instant::now(); + + // Phase 1/1b/1c: GPU data upload + experience collector init (once) + self.init_gpu_data(training_data).await?; + self.init_gpu_raw_buffers(training_data).await?; + self.init_gpu_experience_collector().await?; + + // Phase 3: GPU experience collection + let gpu_experiences_collected = self.collect_gpu_experiences( + training_data, + ).await?; + + // Free DqnGpuData when GPU experience collector is active + if gpu_experiences_collected && self.gpu_data.is_some() { + info!("GPU experience collector active — releasing DqnGpuData to reclaim VRAM"); + self.gpu_data = None; + if let Some(pool) = self.buffer_pool.take() { + drop(pool); + } + } + + // CUDA builds: GPU experience collector is MANDATORY + if !gpu_experiences_collected { + return Err(anyhow::anyhow!( + "GPU experience collector MUST be active for CUDA training. \ + No CPU fallback path exists in CUDA builds. \ + Set enable_gpu_experience_collector=true \ + or check GPU collector initialization errors above." + )); + } + + // Phase 2: Batched training from replay buffer + let train_step_count = self.run_training_steps(training_data).await?; + + // Epoch boundary readback + safety checks + let boundary = if train_step_count > 0 { + let b = self.process_epoch_boundary(epoch, train_step_count, &mut monitor).await?; + Some(b) + } else { + None + }; + + // Sync GPU weights after training + self.sync_gpu_weights().await?; + + // Flush GPU-accumulated max priority + if train_step_count > 0 { + let agent = self.agent.read().await; + if let Err(e) = agent.flush_max_priority() { + debug!("GPU max_priority flush failed (non-fatal): {}", e); + } + } + + let epoch_duration = epoch_start.elapsed(); + + // Calculate and log epoch metrics + let log_output = self.log_epoch_metrics_and_financials( + epoch, + train_step_count, + &boundary, + &mut monitor, + epoch_duration, + &mut total_action_counts, + &mut total_factored_action_counts, + ).await?; + + total_loss += log_output.avg_loss; + total_q_value += log_output.avg_q_value; + total_gradient_norm += log_output.avg_grad_norm; + + let epoch_avg_reward = if !monitor.reward_history.is_empty() { + monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 + } else { + 0.0 + }; + total_reward += epoch_avg_reward as f64; + + // Checkpoints + early stopping (returns Err to signal hyperopt on stop) + self.handle_epoch_checkpoints_and_early_stopping( + epoch, + train_step_count, + &log_output, + &mut checkpoint_callback, + ).await?; + + // WAVE 13-A2: Save periodic checkpoint every N epochs + if (epoch + 1) % self.hyperparams.checkpoint_frequency == 0 { + info!( + "Saving periodic checkpoint at epoch {}/{}", + epoch + 1, self.hyperparams.epochs + ); + let checkpoint_data = self.serialize_model().await?; + let checkpoint_size = checkpoint_data.len(); + let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false) + .context("Failed to save periodic checkpoint")?; + info!( + "Periodic checkpoint saved: {} ({} bytes)", + checkpoint_path, checkpoint_size + ); + } + } + + let training_duration = start_time.elapsed(); + + let metrics = self.create_final_metrics( + total_loss, total_q_value, total_gradient_norm, total_reward, + self.hyperparams.epochs, training_duration, false, + total_action_counts, total_factored_action_counts, + ).await?; + + { + let mut stored_metrics = self.metrics.write().await; + *stored_metrics = metrics.clone(); + } + + info!( + "Training completed in {:.2}s: final_loss={:.6}, avg_q_value={:.4}", + training_duration.as_secs_f64(), + metrics.loss, + metrics.additional_metrics.get("avg_q_value").unwrap_or(&0.0) + ); + + info!("Best model summary:"); + info!( + " Best Sharpe: {:.4} at epoch {} (val_loss={:.6})", + self.best_sharpe, self.best_epoch, self.best_val_loss + ); + info!(" Best model checkpoint: best_model.safetensors"); + + Ok(metrics) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: one-time training configuration logging + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn log_training_config(&mut self) { + match self.hyperparams.target_update_mode { + TargetUpdateMode::Soft => { + let half_life = convergence_half_life(self.hyperparams.tau); + info!("WAVE 16: Using soft target updates (Polyak averaging)"); + info!(" Tau: {}", self.hyperparams.tau); + info!(" Convergence half-life: {} steps", half_life as usize); + info!(" Strategy: Smooth Q-value tracking (50-70% variance reduction)"); + }, + TargetUpdateMode::Hard => { + info!("WAVE 16: Using hard target updates (legacy mode)"); + info!(" Update frequency: every 1000 steps"); + info!(" Warning: Sudden Q-value shifts may cause instability"); + }, + } + + info!("Rainbow DQN Components:"); + info!(" Double DQN (always enabled)"); + + if self.hyperparams.use_dueling { + info!(" Dueling Networks (value/advantage streams, hidden_dim={})", self.hyperparams.dueling_hidden_dim); + } else { + info!(" Dueling Networks: disabled"); + } + + if self.hyperparams.use_per { + info!(" Prioritized Experience Replay (a={}, b={}->1.0)", + self.hyperparams.per_alpha, self.hyperparams.per_beta_start); + } else { + info!(" Prioritized Experience Replay: disabled"); + } + + if self.hyperparams.n_steps > 1 { + info!(" N-Step Returns (n={})", self.hyperparams.n_steps); + } else { + info!(" N-Step Returns (n=1, standard TD)"); + } + + if self.hyperparams.use_distributional { + info!(" Categorical DQN (atoms={}, V=[{}, {}])", + self.hyperparams.num_atoms, self.hyperparams.v_min, self.hyperparams.v_max); + } else { + info!(" Categorical DQN / C51: disabled"); + } + + if self.hyperparams.use_noisy_nets { + info!(" Noisy Networks (sigma_init={})", self.hyperparams.noisy_sigma_init); + } else { + info!(" Noisy Networks: disabled"); + } + + if self.hyperparams.use_noisy_nets { + let floor = self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05); + let mut agent = self.agent.write().await; + agent.set_epsilon(floor); + info!(" Noisy nets active: epsilon set to noisy_epsilon_floor={:.4}", floor); + } + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: reset per-epoch state + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) fn reset_epoch_state(&mut self, _epoch: usize) { + self.metrics_aggregator.reset(); + self.pnl_history.clear(); + self.reward_fn.reset_epoch_state(); + + if self.hyperparams.use_dsr { + self.portfolio_tracker.reset(); + } + + // GPU-persistent epoch state: set reset flags instead of CPU state mutation. + if let Some(ref mut collector) = self.gpu_experience_collector { + let mut flags: u32 = 0; + if self.hyperparams.use_dsr { + flags |= 1; // reset portfolio + flags |= 2; // reset DSR normalizer + } + collector.set_reset_flags(flags); + } + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Phase 1 — GPU data upload (once) + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn init_gpu_data( + &mut self, + training_data: &[(FeatureVector, Vec)], + ) -> Result<()> { + if self.gpu_data.is_some() { + return Ok(()); + } + + // Check double-buffer: skip upload if active slot already populated + let skip_upload = self.double_buffer.as_ref().is_some_and(|db| db.active().is_some()); + + if skip_upload { + info!("DoubleBuffer: active slot populated, skipping re-upload"); + return Ok(()); + } + + let upload_result = if let Some(ref mut pool) = self.buffer_pool { + info!("GpuBufferPool: reusing pre-allocated staging buffers for {} bars", training_data.len()); + pool.upload_dqn(training_data, &self.device) + } else { + DqnGpuData::upload(training_data, &self.device) + }; + + match upload_result { + Ok(mut gpu_data) => { + info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)", + gpu_data.num_bars, + gpu_data.feature_dim, + (gpu_data.num_bars * (42 + 4) * 4) as f64 / 1_048_576.0 + ); + // Upload OFI features to GPU if available + if let Some(ref ofi) = self.ofi_features { + match gpu_data.upload_ofi(ofi, &self.device) { + Ok(()) => info!("GPU OFI features uploaded: {} bars x 8 dims", ofi.len()), + Err(e) => { + return Err(anyhow::anyhow!("GPU OFI upload FAILED (no CPU fallback): {e}")); + } + } + } + // Set tensor core alignment so build_*_states pads output + let ofi_enabled = self.ofi_features.is_some(); + let raw_dim = if ofi_enabled { 53 } else { 45 }; + let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device); + gpu_data.set_aligned_state_dim(aligned_dim); + self.gpu_data = Some(gpu_data); + } + Err(e) => { + return Err(anyhow::anyhow!("GPU data pre-upload FAILED (no CPU fallback): {e}")); + } + } + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Phase 1b — raw cudarc targets + features + portfolio sim (once) + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn init_gpu_raw_buffers( + &mut self, + training_data: &[(FeatureVector, Vec)], + ) -> Result<()> { + if self.targets_raw_cuda.is_some() { + return Ok(()); + } + + let candle_core::Device::Cuda(ref cuda_dev) = self.device else { + return Ok(()); + }; + + let stream = cuda_dev.cuda_stream(); + let target_dim = 4; + let feature_dim = 42; + let num_bars = training_data.len(); + + // Build flat targets (same as DqnGpuData::upload but for cudarc) + let mut flat_targets = Vec::with_capacity(num_bars * target_dim); + for (_, targets) in training_data { + for i in 0..target_dim { + flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32); + } + } + + match stream.memcpy_stod(&flat_targets) { + Ok(buf) => { + info!("CUDA targets_raw uploaded: {} bars x 4 ({:.1} KB)", + num_bars, (num_bars * target_dim * 4) as f64 / 1024.0); + self.targets_raw_cuda = Some(buf); + } + Err(e) => { + return Err(anyhow::anyhow!("CUDA targets_raw upload FAILED (no CPU fallback): {e}")); + } + } + + // Build flat features [num_bars * 42] for GPU experience kernel + if self.features_raw_cuda.is_none() { + let mut flat_features = Vec::with_capacity(num_bars * feature_dim); + for (features, _) in training_data { + for &v in features.iter() { + flat_features.push(v as f32); + } + } + match stream.memcpy_stod(&flat_features) { + Ok(buf) => { + info!("CUDA features_raw uploaded: {} bars x {} ({:.1} KB)", + num_bars, feature_dim, (num_bars * feature_dim * 4) as f64 / 1024.0); + self.features_raw_cuda = Some(buf); + } + Err(e) => { + return Err(anyhow::anyhow!("CUDA features_raw upload FAILED (no CPU fallback): {e}")); + } + } + } + + // Initialize GPU portfolio simulator + if self.gpu_portfolio_sim.is_none() { + if self.targets_raw_cuda.is_some() { + use crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator; + match GpuPortfolioSimulator::new( + stream, + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + self.max_position as f32, + EPISODE_LENGTH, + training_data.len(), + ) { + Ok(sim) => { + info!("GPU portfolio simulator initialized"); + self.gpu_portfolio_sim = Some(sim); + } + Err(e) => { + return Err(anyhow::anyhow!("GPU portfolio sim init FAILED (no CPU fallback): {e}")); + } + } + } + } + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Phase 1c — GPU experience collector init (once) + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn init_gpu_experience_collector(&mut self) -> Result<()> { + if !self.hyperparams.enable_gpu_experience_collector + || self.gpu_experience_collector.is_some() + { + return Ok(()); + } + + let candle_core::Device::Cuda(ref cuda_dev) = self.device else { + return Ok(()); + }; + + use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector; + let stream = cuda_dev.cuda_stream(); + + // Read-lock agent to extract dueling network weights for GPU collector. + let agent = self.agent.read().await; + let dqn_ref: Option<&crate::dqn::DQN> = match &*agent { + DQNAgentType::Standard(ref dqn) => Some(dqn), + DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()), + }; + + let init_result = if let Some(dqn) = dqn_ref { + let curiosity_vars = self.curiosity_module.as_ref() + .map(|c| c.forward_model_vars()); + + let state_dim = agent.get_state_dim(); + let market_dim: usize = 42; + let num_atoms_max = (self.hyperparams.num_atoms as usize).max(51); + let kernel_dims = (state_dim, market_dim, num_atoms_max); + + let use_branching = self.hyperparams.use_branching; + + // Priority: branching > plain dueling > hybrid (distributional+dueling) + if let (true, Some(online_br), Some(target_br)) = ( + use_branching, + dqn.branching_q_network.as_ref(), + dqn.branching_target_network.as_ref(), + ) { + let cfg = online_br.config(); + let dims = ( + *cfg.shared_hidden_dims.first().unwrap_or(&256), + *cfg.shared_hidden_dims.get(1).unwrap_or(&256), + cfg.value_hidden_dim, + cfg.branch_hidden_dim, + ); + let alloc_episodes = self.compute_alloc_episodes(state_dim); + Some(GpuExperienceCollector::new( + stream, + online_br.vars(), + target_br.vars(), + curiosity_vars, + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + dims, + kernel_dims, + alloc_episodes, + self.hyperparams.gpu_timesteps_per_episode, + true, // use_branching + )) + } else if let (Some(online), Some(target)) = ( + dqn.dueling_q_network.as_ref(), + dqn.dueling_target_network.as_ref(), + ) { + let cfg = online.config(); + let dims = ( + *cfg.shared_hidden_dims.first().unwrap_or(&256), + *cfg.shared_hidden_dims.get(1).unwrap_or(&256), + cfg.value_hidden_dim, + cfg.advantage_hidden_dim, + ); + let alloc_episodes = self.compute_alloc_episodes(state_dim); + Some(GpuExperienceCollector::new( + stream, + online.vars(), + target.vars(), + curiosity_vars, + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + dims, + kernel_dims, + alloc_episodes, + self.hyperparams.gpu_timesteps_per_episode, + self.hyperparams.use_branching, + )) + } else if let (Some(online), Some(target)) = ( + dqn.dist_dueling_q_network.as_ref(), + dqn.dist_dueling_target_network.as_ref(), + ) { + let cfg = online.config(); + let dims = ( + *cfg.shared_hidden_dims.first().unwrap_or(&256), + *cfg.shared_hidden_dims.get(1).unwrap_or(&256), + cfg.value_hidden_dim, + cfg.advantage_hidden_dim, + ); + let alloc_episodes = self.compute_alloc_episodes(state_dim); + Some(GpuExperienceCollector::new( + stream, + online.vars(), + target.vars(), + curiosity_vars, + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + dims, + kernel_dims, + alloc_episodes, + self.hyperparams.gpu_timesteps_per_episode, + self.hyperparams.use_branching, + )) + } else { + return Err(anyhow::anyhow!( + "GPU experience collector FAILED: no dueling, hybrid, or branching networks — \ + DQN agent has no Q-network. This is a configuration bug." + )); + } + } else { + None + }; + drop(agent); + + if let Some(result) = init_result { + match result { + Ok(mut collector) => { + // Lazy-init GPU monitoring reducer on same stream as collector + if self.gpu_monitoring.is_none() { + match crate::cuda_pipeline::gpu_monitoring::GpuMonitoringReducer::new(collector.stream()) { + Ok(mon) => { + info!("GPU monitoring reducer initialized"); + self.gpu_monitoring = Some(mon); + } + Err(e) => { + return Err(anyhow::anyhow!( + "GPU monitoring reducer init FAILED: {e} — metrics pipeline must be GPU-resident" + )); + } + } + } + // Upload OFI features to GPU if available. + if let Some(ref ofi) = self.ofi_features { + let flat: Vec = ofi.iter() + .flat_map(|f| f.iter().map(|&v| v as f32)) + .collect(); + collector.upload_ofi_features(&flat)?; + } + self.gpu_experience_collector = Some(collector); + } + Err(e) => { + return Err(anyhow::anyhow!( + "GPU experience collector init FAILED (no CPU fallback): {e}" + )); + } + } + } + Ok(()) + } + + /// Compute allocation episode count for GPU buffers. + /// Auto-scales when configured >= 128 (production), respects smaller test overrides. + fn compute_alloc_episodes(&self, state_dim: usize) -> usize { + use ml_core::memory_optimization::detect_gpu_hardware; + let configured = self.hyperparams.gpu_n_episodes; + if configured >= 128 { + match detect_gpu_hardware() { + Ok(hw) => configured.max(hw.optimal_n_episodes( + state_dim, + self.hyperparams.gpu_timesteps_per_episode, + )).min(0x8000), + Err(_) => configured, + } + } else { + configured + } + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Phase 3 — GPU experience collection + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn collect_gpu_experiences( + &mut self, + training_data: &[(FeatureVector, Vec)], + ) -> Result { + // Feature normalization stats calculation epoch (kept for API compat) + let _stats_collection_epochs = { + let ratio_based = (self.hyperparams.epochs as f32 * self.hyperparams.feature_stats_collection_ratio) as usize; + let capped = match self.hyperparams.max_feature_stats_epochs { + Some(max_epochs) => ratio_based.min(max_epochs), + None => ratio_based, + }; + capped.max(1) + }; + + let ( + Some(ref mut collector), + Some(ref features_buf), + Some(ref targets_buf), + ) = ( + &mut self.gpu_experience_collector, + &self.features_raw_cuda, + &self.targets_raw_cuda, + ) else { + return Ok(false); + }; + + use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig; + + let raw_sd = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; + let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &self.device); + + // Cache n_episodes on first epoch + let n_episodes = if let Some(cached) = self.cached_n_episodes { + cached + } else { + use ml_core::memory_optimization::detect_gpu_hardware; + let configured = self.hyperparams.gpu_n_episodes; + let computed = if configured >= 128 { + match detect_gpu_hardware() { + Ok(hw) => { + let optimal = hw.optimal_n_episodes( + aligned_sd, + self.hyperparams.gpu_timesteps_per_episode, + ); + let chosen = configured.max(optimal).min(4096); + if chosen != configured { + info!( + "GPU auto-scaled n_episodes: {} -> {} (SMs={}, VRAM={:.0}MB)", + configured, chosen, hw.sm_count, hw.free_memory_mb + ); + } + chosen as i32 + } + Err(_) => configured as i32, + } + } else { + configured as i32 + }; + self.cached_n_episodes = Some(computed); + computed + }; + + let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32; + let total_bars = training_data.len() as i32; + let usable_bars = (total_bars - timesteps).max(1); + let stride = (usable_bars / n_episodes).max(1); + let episode_starts: Vec = (0..n_episodes) + .map(|i| (i * stride).rem_euclid(usable_bars)) + .collect(); + + // Reset per-episode state before each epoch + if let Err(e) = collector.reset_episodes( + self.hyperparams.initial_capital as f32, + self.hyperparams.avg_spread as f32, + self.hyperparams.cash_reserve_percent as f32, + ) { + return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}")); + } + + let agent = self.agent.read().await; + let epsilon = agent.get_epsilon(); + drop(agent); + + let config = ExperienceCollectorConfig { + n_episodes, + timesteps_per_episode: timesteps, + total_bars, + epsilon, + gamma: self.hyperparams.gamma as f32, + max_position: self.max_position as f32, + enable_action_masking: self.enable_action_masking, + curiosity_scale: if self.curiosity_module.is_some() { 1.0 } else { 0.0 }, + hold_reward: self.hyperparams.hold_penalty.abs() as f32, + tx_cost_multiplier: self.hyperparams.transaction_cost_multiplier as f32, + count_bonus_coefficient: self.hyperparams.count_bonus_coefficient + .unwrap_or(0.0) as f32, + q_clip_min: -500.0, + q_clip_max: 500.0, + huber_kappa: if self.hyperparams.use_huber_loss { + self.hyperparams.huber_delta as f32 + } else { + 0.0 + }, + use_noisy_nets: self.hyperparams.use_noisy_nets, + noisy_sigma_init: self.hyperparams.noisy_sigma_init as f32, + use_distributional: self.hyperparams.use_distributional, + num_atoms: self.hyperparams.num_atoms as i32, + v_min: self.hyperparams.v_min as f32, + v_max: self.hyperparams.v_max as f32, + fill_median_spread: self.hyperparams.avg_spread as f32, + fill_median_vol: self.median_vol as f32, + fill_ioc_fill_prob: 0.85, + fill_limit_fill_min: 0.30, + fill_limit_fill_max: 0.80, + fill_spread_cost_frac: 0.50, + fill_spread_capture_frac: 0.50, + fill_simulation_enabled: self.median_vol > 0.0, + use_dsr: self.hyperparams.use_dsr, + dsr_eta: self.hyperparams.dsr_eta as f32, + n_steps: self.hyperparams.n_steps as i32, + ..Default::default() + }; + + // Zero-roundtrip GPU path — GPU PER is always active in CUDA builds + let gpu_batch = collector.collect_experiences_gpu( + features_buf, targets_buf, &episode_starts, &config, &self.device, + ).map_err(|e| anyhow::anyhow!( + "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" + ))?; + + let count = gpu_batch.n_episodes * gpu_batch.timesteps; + info!("GPU collected {} experiences (zero-roundtrip, {} episodes x {} timesteps)", + count, gpu_batch.n_episodes, gpu_batch.timesteps); + + if let Some(ref mut mon) = self.gpu_monitoring { + if let Err(e) = mon.reduce(collector.rewards_gpu(), collector.actions_gpu(), count) { + debug!("GPU monitoring reduce failed (non-fatal): {e}"); + } + } + + if count > 0 { + if let Err(e) = collector.train_curiosity_gpu( + gpu_batch.n_episodes, gpu_batch.timesteps, + ) { + debug!("GPU curiosity training failed (non-fatal): {e}"); + } + } + + if count > 0 { + let agent = self.agent.read().await; + agent.insert_batch_tensors( + &gpu_batch.states, + &gpu_batch.next_states, + &gpu_batch.actions, + &gpu_batch.rewards, + &gpu_batch.dones, + ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; + } + Ok(true) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Batched training steps with guard kernel + // ═══════════════════════════════════════════════════════════════════════ + + /// Runs all training steps for one epoch, returning the step count. + pub(crate) async fn run_training_steps( + &mut self, + training_data: &[(FeatureVector, Vec)], + ) -> Result { + let batch_size = self.hyperparams.batch_size; + let num_training_steps = if self.can_train().await? { + let full_steps = (training_data.len() / batch_size).max(1); + let cap = self.hyperparams.max_training_steps_per_epoch; + if cap > 0 { full_steps.min(cap) } else { full_steps } + } else { + 0 + }; + + let mut train_step_count = 0; + + // Lazy-init fused CUDA training context (Standard DQN only). + // Recreate if batch_size changed (OOM recovery). + if num_training_steps > 0 { + self.ensure_fused_ctx().await; + } + + // Lazy-init + reset guard accumulators for this epoch + { + if self.training_guard.is_none() && self.device.is_cuda() { + match crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard::new(&self.device) { + Ok(g) => { + info!("GPU training guard initialized (epoch loop)"); + self.training_guard = Some(g); + } + Err(e) => return Err(anyhow::anyhow!("GPU training guard init: {e}")), + } + } + if let Some(ref mut guard) = self.training_guard { + guard.reset_accumulators() + .map_err(|e| anyhow::anyhow!("guard reset: {e}"))?; + } + } + + let guard_collapse_thresh = + self.hyperparams.learning_rate as f32 + * self.hyperparams.gradient_collapse_multiplier as f32; + let guard_past_warmup = { + let ws = (self.collapse_warmup_buffer_size as f64 * 0.2) as u64; + self.gradient_logging_step as u64 > ws + }; + + // Batch pre-sampling: sample K batches under one READ lock + const PREFETCH_K: usize = 32; + for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) { + let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps); + + let batches = { + let agent = self.agent.read().await; + let buffer = agent.memory(); + let mut b = Vec::with_capacity(chunk_end - chunk_start); + for _ in chunk_start..chunk_end { + b.push(buffer.can_sample(self.current_batch_size).then(|| { + buffer + .sample(self.current_batch_size) + .map_err(|e| anyhow::anyhow!("Pre-sample: {e}")) + }).transpose()?); + } + b + }; + + // GPU train steps (single WRITE lock) + { + let mut agent = self.agent.write().await; + let accum_steps = self.hyperparams.gradient_accumulation_steps; + + if accum_steps <= 1 { + for explicit_batch in batches { + // Fused CUDA path: 3 kernels + CUDA Graph, no Candle dispatch + let _gpu_result = if let Some(ref mut fused) = self.fused_ctx { + let batch_data = explicit_batch.as_ref().ok_or_else(|| { + anyhow::anyhow!("No batch data for fused training step") + })?; + fused.run_full_step(batch_data, &mut *agent, &self.device)? + } else { + agent.train_step(explicit_batch) + .map_err(|e| anyhow::anyhow!("Train step FAILED: {e}"))? + }; + + if let Some(ref mut guard) = self.training_guard { + let gr = guard.check_and_accumulate( + &_gpu_result.loss_gpu, + &_gpu_result.grad_norm_gpu, + 1e6_f32, + guard_collapse_thresh, + !guard_past_warmup, + ).map_err(|e| anyhow::anyhow!("guard check: {e}"))?; + if gr.halt_nan { + return Err(anyhow::anyhow!( + "NaN/Inf at step {}: loss={}, grad={}", + train_step_count, gr.raw_loss, gr.raw_grad_norm + )); + } + if gr.halt_grad_collapse { + agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + } + train_step_count += 1; + self.gradient_logging_step += 1; + } + } else { + // Gradient accumulation path + let mut batch_iter = batches.into_iter().peekable(); + + while batch_iter.peek().is_some() { + let mut accumulated_grads: Option = None; + let mut group_td_gpu: Vec = Vec::new(); + let mut group_idx_gpu: Vec = Vec::new(); + let mut accum_count: usize = 0; + + for _ in 0..accum_steps { + let batch = match batch_iter.next() { + Some(b) => b, + None => break, + }; + + let result = match agent.compute_gradients(batch) { + Ok(r) => r, + Err(e) => { + let msg = e.to_string(); + if msg.contains("Early stopping") || msg.contains("Gradient collapse") { + return Err(anyhow::anyhow!("{}", msg)); + } + return Err(anyhow::anyhow!("GPU gradient compute FAILED: {e}")); + } + }; + + let vars = agent.optimizer_vars() + .map_err(|e| anyhow::anyhow!("optimizer vars: {e}"))?; + crate::gradient_accumulation::accumulate_grads( + &mut accumulated_grads, + result.grads, + &vars, + ).map_err(|e| anyhow::anyhow!("grad accum: {e}"))?; + + { + if let (Some(ref loss_t), Some(ref gn_t)) = + (&result.loss_tensor_gpu, &result.grad_norm_gpu) + { + if let Some(ref mut guard) = self.training_guard { + let gr = guard.check_and_accumulate( + loss_t, gn_t, 1e6_f32, + guard_collapse_thresh, !guard_past_warmup, + ).map_err(|e| anyhow::anyhow!("guard accum step: {e}"))?; + if gr.halt_nan { + return Err(anyhow::anyhow!( + "NaN/Inf in accum sub-step: loss={}, grad={}", + gr.raw_loss, gr.raw_grad_norm + )); + } + if gr.halt_grad_collapse { + agent.check_gradient_collapse(gr.raw_grad_norm).map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + } + } + if let Some(td) = result.td_errors_gpu { + group_td_gpu.push(td); + } + if let Some(idx) = result.indices_gpu { + group_idx_gpu.push(idx); + } + } + + accum_count += 1; + } + + if accum_count == 0 { break; } + + // Scale and apply (single optimizer step) + if let Some(ref mut grads) = accumulated_grads { + let vars = agent.optimizer_vars() + .map_err(|e| anyhow::anyhow!("optimizer vars: {e}"))?; + crate::gradient_accumulation::scale_grads( + grads, &vars, 1.0 / accum_count as f64, + ).map_err(|e| anyhow::anyhow!("grad scale: {e}"))?; + + agent.apply_accumulated_gradients(grads) + .map_err(|e| anyhow::anyhow!("apply grads: {e}"))?; + } + + // Update priorities for all batches in this group + if !group_td_gpu.is_empty() && !group_idx_gpu.is_empty() { + let td_cat = Tensor::cat(&group_td_gpu, 0) + .map_err(|e| anyhow::anyhow!("TD cat: {e}"))?; + let idx_cat = Tensor::cat(&group_idx_gpu, 0) + .map_err(|e| anyhow::anyhow!("idx cat: {e}"))?; + agent.update_priorities_gpu(&idx_cat, &td_cat) + .map_err(|e| anyhow::anyhow!("PER update: {e}"))?; + } + agent.step_replay_buffer(); + + train_step_count += accum_count; + self.gradient_logging_step += accum_count; + } + } + } + } + + Ok(train_step_count) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Epoch boundary readback + safety checks + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn process_epoch_boundary( + &mut self, + epoch: usize, + train_step_count: usize, + monitor: &mut TrainingMonitor, + ) -> Result { + let (avg_loss, avg_grad) = if let Some(ref mut guard) = self.training_guard { + let (al, ag) = guard.read_accumulators() + .map_err(|e| anyhow::anyhow!("guard read_accumulators: {e}"))?; + (al as f32, ag as f32) + } else { + (0.0_f32, 0.0_f32) + }; + + let avg_q = { + let mut agent = self.agent.write().await; + self.estimate_avg_q_value_with_early_stopping(&mut agent).await? + }; + + // Epoch-level safety checks + if !avg_loss.is_finite() || !avg_grad.is_finite() { + let msg = format!( + "SAFETY: NaN/Inf in epoch {} avg — loss={:.6}, grad={:.6}", + epoch, avg_loss, avg_grad + ); + match self.safety_level { + crate::safety::SafetyLevel::Strict => { + return Err(anyhow::anyhow!("{} (stopping training)", msg)); + }, + crate::safety::SafetyLevel::Normal | crate::safety::SafetyLevel::Permissive => { + debug!("{} (continuing training)", msg); + }, + } + } + + // Loss history (epoch-level average) + self.safety_loss_history.push_back(avg_loss); + if self.safety_loss_history.len() > 30 { + self.safety_loss_history.pop_front(); + } + + let n = train_step_count as f64; + + // Metrics aggregator + let current_avg_reward = if !monitor.reward_history.is_empty() { + monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 + } else { + 0.0 + }; + self.metrics_aggregator.record(avg_loss, avg_q as f32, current_avg_reward); + self.metrics_aggregator.record_gradient(avg_grad); + + if self.metrics_aggregator.should_log(&self.logging_config) { + let aggregated = self.metrics_aggregator.aggregate_and_clear(); + log_training_progress(&aggregated); + } + + // Gradient diagnostics at epoch boundary + { + let mut agent = self.agent.write().await; + agent.log_diagnostics(avg_grad) + .map_err(|e| { + tracing::info!("Early stopping (gradient collapse): {}", e); + anyhow::anyhow!("Early stopping: {}", e) + })?; + } + + monitor.track_q_value_range(avg_q); + + // Return raw accumulated totals for epoch metric computation + // The caller computes per-step averages from epoch_loss/n, etc. + Ok(EpochBoundaryMetrics { + avg_loss: (avg_loss as f64 * n) as f32, + avg_grad: (avg_grad as f64 * n) as f32, + avg_q: avg_q * n, + }) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Sync GPU weight copies after training updates + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn sync_gpu_weights(&mut self) -> Result<()> { + let Some(ref mut collector) = self.gpu_experience_collector else { + return Ok(()); + }; + + let agent = self.agent.read().await; + let dqn_ref: Option<&crate::dqn::DQN> = match &*agent { + DQNAgentType::Standard(ref dqn) => Some(dqn), + DQNAgentType::RegimeConditional(ref regime_dqn) => Some(regime_dqn.primary_head()), + }; + + if let Some(dqn) = dqn_ref { + // Sync online weights: branching > plain dueling > hybrid + let online_synced = if let Some(ref bn) = dqn.branching_q_network { + collector.sync_online_weights(bn.vars()).is_ok() + } else if let Some(ref online) = dqn.dueling_q_network { + collector.sync_online_weights(online.vars()).is_ok() + } else if let Some(ref online) = dqn.dist_dueling_q_network { + collector.sync_online_weights(online.vars()).is_ok() + } else { + false + }; + if !online_synced { + return Err(anyhow::anyhow!("GPU online weight sync FAILED -- stale Q-values would corrupt training")); + } + + // Sync target weights: branching > plain dueling > hybrid + let target_synced = if let Some(ref bn) = dqn.branching_target_network { + collector.sync_target_weights(bn.vars()).is_ok() + } else if let Some(ref target) = dqn.dueling_target_network { + collector.sync_target_weights(target.vars()).is_ok() + } else if let Some(ref target) = dqn.dist_dueling_target_network { + collector.sync_target_weights(target.vars()).is_ok() + } else { + false + }; + if !target_synced { + return Err(anyhow::anyhow!("GPU target weight sync FAILED -- stale target Q-values would corrupt training")); + } + + // Sync branching DQN extra heads (order + urgency, branches 1+2) + if let Some(ref bn) = dqn.branching_q_network { + collector.sync_online_branching(bn.vars()) + .map_err(|e| anyhow::anyhow!("GPU online branching weight sync FAILED: {e}"))?; + } + if let Some(ref bn) = dqn.branching_target_network { + collector.sync_target_branching(bn.vars()) + .map_err(|e| anyhow::anyhow!("GPU target branching weight sync FAILED: {e}"))?; + } + + // Sync RMSNorm weights for distributional dueling networks (D6) + if let Some(ref online) = dqn.dist_dueling_q_network { + collector.sync_online_rmsnorm(online.vars()) + .map_err(|e| anyhow::anyhow!("GPU online RMSNorm weight sync FAILED: {e}"))?; + } + if let Some(ref target) = dqn.dist_dueling_target_network { + collector.sync_target_rmsnorm(target.vars()) + .map_err(|e| anyhow::anyhow!("GPU target RMSNorm weight sync FAILED: {e}"))?; + } + } + drop(agent); + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: M2 — Refresh stale PER priorities + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn refresh_stale_per_priorities(&self, epoch: usize) -> Result<()> { + if !self.hyperparams.use_per || self.hyperparams.n_steps <= 1 { + return Ok(()); + } + + let agent = self.agent.read().await; + let buffer = agent.memory(); + let buf_len = buffer.len(); + let refresh_limit = (buf_len / 20).max(32).min(256); + let max_age = 500; + let stale_indices = buffer.get_stale_indices(max_age, refresh_limit); + + if stale_indices.is_empty() { + return Ok(()); + } + + let experiences = buffer.get_experiences_at(&stale_indices); + + let mut flat_states: Vec = Vec::with_capacity(stale_indices.len() * 64); + let mut valid_indices = Vec::with_capacity(stale_indices.len()); + let mut state_dim: Option = None; + + for (idx, exp_opt) in stale_indices.iter().zip(experiences.iter()) { + if let Some(Some(exp)) = exp_opt.as_ref().map(Some) { + if let Some(sd) = state_dim { + if exp.state.len() != sd { + continue; + } + } else { + state_dim = Some(exp.state.len()); + } + flat_states.extend_from_slice(&exp.state); + valid_indices.push(*idx); + } + } + + if let Some(sd) = state_dim { + if !valid_indices.is_empty() { + let batch_size_refresh = valid_indices.len(); + let td_result = candle_core::Tensor::from_vec( + flat_states, &[batch_size_refresh, sd], agent.device(), + ) + .and_then(|bt| { + agent.forward(&bt).map_err(|e| { + candle_core::Error::Msg(format!("forward: {e}")) + }) + }) + .and_then(|q_vals| q_vals.max(candle_core::D::Minus1)) + .and_then(|mq| { + let abs_q = mq.abs()?; + let floor = candle_core::Tensor::new(0.01_f32, abs_q.device())?; + let clamped = abs_q.broadcast_maximum(&floor)?; + let nan_mask = clamped.ne(&clamped)?; + let floor_bcast = floor.broadcast_as(clamped.shape())?; + let safe = nan_mask.where_cond(&floor_bcast, &clamped)?; + Ok(safe) + }); + + if let Ok(td_errors_gpu) = td_result { + let idx_u32: Vec = valid_indices.iter() + .map(|&i| i as u32) + .collect(); + if let Ok(idx_tensor) = candle_core::Tensor::new( + idx_u32, agent.device(), + ) { + if let Err(e) = agent.update_priorities_gpu( + &idx_tensor, &td_errors_gpu, + ) { + debug!("M2: GPU priority refresh failed (non-fatal): {}", e); + } else { + debug!( + "Epoch {}: Refreshed {} stale PER priorities (max_age={}, GPU)", + epoch + 1, valid_indices.len(), max_age + ); + } + } + } + } + } + + Ok(()) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Epoch metrics logging + financials + Prometheus + QuestDB + // ═══════════════════════════════════════════════════════════════════════ + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn log_epoch_metrics_and_financials( + &mut self, + epoch: usize, + train_step_count: usize, + boundary: &Option, + monitor: &mut TrainingMonitor, + epoch_duration: std::time::Duration, + total_action_counts: &mut [usize; 5], + total_factored_action_counts: &mut [usize; 45], + ) -> Result { + // Calculate epoch metrics (average over training steps) + let (epoch_loss, epoch_q_value, epoch_gradient_norm) = match boundary { + Some(b) => (b.avg_loss as f64, b.avg_q, b.avg_grad as f64), + None => (0.0, 0.0, 0.0), + }; + + let (avg_loss, avg_q_value, avg_grad_norm) = if train_step_count > 0 { + ( + epoch_loss / train_step_count as f64, + epoch_q_value / train_step_count as f64, + epoch_gradient_norm / train_step_count as f64, + ) + } else { + (0.0, 0.0, 0.0) + }; + + // WAVE 30: Log epoch end + use crate::dqn::logging::AggregatedMetrics; + let epoch_avg_reward = if !monitor.reward_history.is_empty() { + monitor.reward_history.iter().sum::() / monitor.reward_history.len() as f32 + } else { + 0.0 + }; + let epoch_metrics = AggregatedMetrics { + mean_loss: avg_loss as f32, + std_loss: 0.0, + mean_q_value: avg_q_value as f32, + std_q_value: 0.0, + mean_reward: epoch_avg_reward, + mean_gradient_norm: avg_grad_norm as f32, + batch_count: train_step_count, + }; + log_epoch_end(epoch + 1, &epoch_metrics, epoch_duration.as_secs_f64()); + + // Download GPU monitoring summary + if let Some(ref mon) = self.gpu_monitoring { + if let Ok(summary) = mon.download_summary() { + if summary.total_experiences > 0 { + info!( + "GPU epoch summary: mean_reward={:.6}, std={:.6}, sharpe={:.3}, actions={:?}", + summary.mean_reward, summary.reward_std, summary.sharpe_estimate, + summary.action_counts + ); + self.pnl_history.push_back(summary.mean_reward as f64); + if self.pnl_history.len() > 1000 { + self.pnl_history.pop_front(); + } + monitor.track_reward(summary.mean_reward); + for (idx, &count) in summary.action_counts.iter().enumerate() { + if count > 0 { + if let Ok(exp_level) = crate::dqn::action_space::ExposureLevel::from_index(idx) { + let factored = self.route_action(exp_level, self.hyperparams.avg_spread as f32); + monitor.track_action(&factored); + } else { + monitor.track_action_by_exposure(idx); + } + } + } + } + } + } + + // Reward statistics every 10 epochs + if (epoch + 1) % 10 == 0 && !monitor.reward_history.is_empty() { + let rewards = &monitor.reward_history; + let reward_mean = rewards.iter().sum::() / rewards.len() as f32; + let reward_variance = rewards.iter() + .map(|r| (r - reward_mean).powi(2)) + .sum::() / rewards.len() as f32; + let reward_std = reward_variance.sqrt(); + let reward_min = rewards.iter().copied().fold(f32::INFINITY, f32::min); + let reward_max = rewards.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let non_zero_count = rewards.iter().filter(|&&r| r.abs() > 1e-9).count(); + let non_zero_pct = (non_zero_count as f32 / rewards.len() as f32) * 100.0; + + info!( + "REWARD_STATS: epoch={}, mean={:.6}, std={:.6}, min={:.6}, max={:.6}, non_zero={}/{} ({:.1}%)", + epoch + 1, reward_mean, reward_std, reward_min, reward_max, + non_zero_count, rewards.len(), non_zero_pct + ); + } + + // Safety plateau detection + if self.safety_loss_history.len() >= 10 { + let recent_losses: Vec = self.safety_loss_history + .iter().rev().take(10).copied().collect(); + + let mean: f32 = recent_losses.iter().sum::() / 10.0; + let variance: f32 = recent_losses.iter() + .map(|&x| (x - mean).powi(2)) + .sum::() / 10.0; + let std_dev = variance.sqrt(); + + if std_dev < mean * 0.01 && mean > 1e-6 { + self.safety_loss_plateau_counter += 1; + if self.safety_loss_plateau_counter >= 10 { + debug!( + "SAFETY: Training stuck for {} epochs (loss variance: {:.6}, mean: {:.6})", + self.safety_loss_plateau_counter, std_dev, mean + ); + } + } else { + self.safety_loss_plateau_counter = 0; + } + } + + // Accumulate action counts + for (i, count) in monitor.action_counts.iter().enumerate() { + total_action_counts[i] += count; + } + for (i, count) in monitor.factored_action_counts.iter().enumerate() { + total_factored_action_counts[i] += count; + } + + // Monitoring validation + if let Err(e) = monitor.validate_all() { + return Err(e); + } + + let current_epsilon = self.get_epsilon().await?; + let (q_min, q_max, q_mean) = monitor.get_q_value_stats(); + + info!( + "Epoch {}/{}: train_loss={:.6}, Q-value={:.4}, grad_norm={:.6}, train_steps={}, epsilon={:.4}, duration={:.2}s", + epoch + 1, self.hyperparams.epochs, + avg_loss, avg_q_value, avg_grad_norm, train_step_count, + current_epsilon, epoch_duration.as_secs_f64() + ); + + // Epsilon decay (skip when noisy nets) + if !self.hyperparams.use_noisy_nets { + let mut agent = self.agent.write().await; + agent.update_epsilon(); + } + + // Noisy sigma schedule + if let Some(ref mut scheduler) = self.noisy_sigma_scheduler { + for _ in 0..train_step_count { + scheduler.step(); + } + let sigma_scale = scheduler.get_sigma(); + let mut agent = self.agent.write().await; + agent.set_noise_sigma_scale(sigma_scale); + } + + // Reset count-based exploration bonus + { + let mut agent = self.agent.write().await; + agent.reset_count_bonus(); + } + + // M2: Refresh stale PER priorities + self.refresh_stale_per_priorities(epoch).await?; + + // LR scheduler + self.lr_scheduler.step(); + let current_lr = self.lr_scheduler.get_lr(); + if epoch % 10 == 0 { + info!( + "Learning rate scheduled update: epoch={}, lr={:.2e} (initial={:.2e})", + epoch + 1, current_lr, self.lr_scheduler.get_initial_lr() + ); + } + let initial_lr = self.lr_scheduler.get_initial_lr(); + if initial_lr > 0.0 { + let decay_factor = current_lr / initial_lr; + self.agent.write().await.update_learning_rate(decay_factor)?; + } + training_metrics::set_learning_rate("dqn", "current", current_lr); + + // Q-value diagnostics + if train_step_count > 0 { + info!( + "Epoch {}/{}: Q-value range=[{:.2}, {:.2}], mean={:.2}", + epoch + 1, self.hyperparams.epochs, q_min, q_max, q_mean + ); + + const Q_VALUE_WARNING_THRESHOLD: f64 = 500_000.0; + if q_max > Q_VALUE_WARNING_THRESHOLD { + warn!( + "Q-value explosion detected at epoch {}: max Q-value {:.2e} exceeds threshold {:.2e}", + epoch + 1, q_max, Q_VALUE_WARNING_THRESHOLD + ); + warn!("Consider:"); + warn!(" Reducing learning rate (current: {:.2e})", self.hyperparams.learning_rate); + warn!(" Enabling target network soft updates (Polyak averaging, tau=0.005)"); + warn!(" Adjusting reward scaling"); + training_metrics::record_gradient_explosion("dqn", "current"); + } + + // M3: Combined Q-value diagnostics + if let Some(((gap_mean, gap_min, gap_max), per_action_avgs)) = + self.compute_epoch_q_diagnostics().await + { + info!( + "Epoch {}/{}: Q-value gap (best-2nd): mean={:.4}, min={:.4}, max={:.4}", + epoch + 1, self.hyperparams.epochs, + gap_mean, gap_min, gap_max + ); + + let names = ["S100", "S50", "Flat", "L50", "L100"]; + let parts: Vec = names.iter() + .zip(per_action_avgs.iter()) + .map(|(n, q)| format!("{}={:.4}", n, q)) + .collect(); + info!( + "Epoch {}/{}: Per-action Q: {}", + epoch + 1, self.hyperparams.epochs, parts.join(", ") + ); + } + } + + // Validation loss + let val_loss = self.compute_validation_loss().await?; + info!( + "Epoch {}/{}: val_loss={} (backtest Sharpe proxy)", + epoch + 1, self.hyperparams.epochs, + if val_loss.abs() < 1e-10 { "N/A".to_owned() } else { format!("{val_loss:.6}") } + ); + + // Prometheus metrics + training_metrics::set_epoch("dqn", "current", (epoch + 1) as f64); + training_metrics::set_epoch_loss("dqn", "current", avg_loss); + training_metrics::set_validation_loss("dqn", "current", val_loss); + if epoch_duration.as_secs_f64() > 0.0 { + training_metrics::set_batches_per_second( + "dqn", "current", + train_step_count as f64 / epoch_duration.as_secs_f64(), + ); + } + training_metrics::set_q_value_stats("dqn", "current", q_mean, q_max); + training_metrics::set_gradient_norm("dqn", "current", avg_grad_norm); + training_metrics::set_epoch_duration("dqn", "current", epoch_duration.as_secs_f64()); + { + let agent = self.agent.read().await; + if let Ok(buf_size) = agent.get_replay_buffer_size() { + training_metrics::set_replay_buffer_size("dqn", "current", buf_size as f64); + } + } + + // Adaptive tau + if epoch > 0 { + let q_mean_growth = q_mean - self.prev_epoch_q_mean; + if q_mean_growth > 0.005 { + self.adaptive_tau = (self.adaptive_tau * 2.0).min(0.01); + warn!( + "Q-value drift detected (delta={:.3}), increasing tau to {:.4}", + q_mean_growth, self.adaptive_tau + ); + } else if q_mean_growth < 0.001 { + self.adaptive_tau = (self.adaptive_tau * 0.9).max(self.hyperparams.tau); + } + { + let mut agent = self.agent.write().await; + match &mut *agent { + DQNAgentType::Standard(dqn) => dqn.config.tau = self.adaptive_tau, + DQNAgentType::RegimeConditional(regime) => { + regime.primary_head_mut().config.tau = self.adaptive_tau; + } + } + } + } + self.prev_epoch_q_mean = q_mean; + + // Action diversity + let epoch_total_factored: usize = monitor.factored_action_counts.iter().sum(); + let epoch_total_exposure: usize = monitor.action_counts.iter().sum(); + let (active_actions_count, action_space_size, diversity_percentage) = + if epoch_total_factored > 0 { + let active_threshold = (epoch_total_factored as f64 * 0.005).max(1.0); + let active = monitor.factored_action_counts.iter() + .filter(|&&count| count as f64 >= active_threshold) + .count(); + (active, 45_usize, (active as f64 / 45.0) * 100.0) + } else { + let active_threshold = (epoch_total_exposure as f64 * 0.005).max(1.0); + let active = monitor.action_counts.iter() + .filter(|&&count| count as f64 >= active_threshold) + .count(); + (active, 5_usize, (active as f64 / 5.0) * 100.0) + }; + + info!( + "Epoch {}/{}: Action diversity={}/{} ({:.1}%)", + epoch + 1, self.hyperparams.epochs, + active_actions_count, action_space_size, diversity_percentage + ); + + // Epoch entropy + let (epoch_entropy, entropy_total, entropy_size) = if epoch_total_factored > 0 { + let mut entropy_raw = 0.0_f64; + for &count in monitor.factored_action_counts.iter() { + if count > 0 { + let p = count as f64 / epoch_total_factored as f64; + entropy_raw -= p * p.ln(); + } + } + (entropy_raw / 45.0_f64.ln(), epoch_total_factored, 45_usize) + } else if epoch_total_exposure > 0 { + let mut entropy_raw = 0.0_f64; + for &count in monitor.action_counts.iter() { + if count > 0 { + let p = count as f64 / epoch_total_exposure as f64; + entropy_raw -= p * p.ln(); + } + } + (entropy_raw / 5.0_f64.ln(), epoch_total_exposure, 5_usize) + } else { + (0.0, 0, 5) + }; + let _ = (entropy_total, entropy_size); + + info!( + " Exploration: entropy={:.3} (1.0=uniform), epsilon={:.4}, noisy_nets={}, count_bonus={}", + epoch_entropy, + self.get_epsilon().await.unwrap_or(0.0), + self.hyperparams.use_noisy_nets, + self.hyperparams.count_bonus_coefficient.unwrap_or(0.0) > 0.0, + ); + + training_metrics::set_action_entropy("dqn", "current", epoch_entropy); + training_metrics::set_action_diversity("dqn", "current", diversity_percentage / 100.0); + + const DIVERSITY_THRESHOLD: usize = 2; + if active_actions_count < DIVERSITY_THRESHOLD { + warn!( + "LOW ACTION DIVERSITY: {}/5 exposure levels (<40%), entropy={:.3}", + active_actions_count, epoch_entropy, + ); + } + + // Episode stats + let (mean_len, std_len, min_len, max_len, exit_counts) = monitor.get_episode_stats(); + if !monitor.episode_lengths.is_empty() { + let total_episodes = monitor.episode_lengths.len(); + + info!( + "WAVE P2 Episode Stats [Epoch {}]: {} episodes, length: mean={:.1}+/-{:.1}, min={}, max={}", + epoch + 1, total_episodes, mean_len, std_len, min_len, max_len + ); + + info!( + " Exit breakdown: profit={}({:.1}%), stop={}({:.1}%), time={}({:.1}%), boundary={}({:.1}%)", + exit_counts[0], (exit_counts[0] as f64 / total_episodes as f64) * 100.0, + exit_counts[1], (exit_counts[1] as f64 / total_episodes as f64) * 100.0, + exit_counts[2], (exit_counts[2] as f64 / total_episodes as f64) * 100.0, + exit_counts[3], (exit_counts[3] as f64 / total_episodes as f64) * 100.0, + ); + } + + // VaR/CVaR + if self.pnl_history.len() > 20 { + let capital = self.hyperparams.initial_capital as f64; + let returns: Vec = if capital > 0.0 { + self.pnl_history.iter().map(|&pnl| pnl / capital).collect() + } else { + self.pnl_history.iter().copied().collect() + }; + + let (var_95, cvar_95) = calculate_var_cvar(&returns, 0.05); + let (var_99, cvar_99) = calculate_var_cvar(&returns, 0.01); + + info!( + "Epoch {}/{}: Risk Metrics - VaR(95%)={:.4}%, CVaR(95%)={:.4}%, VaR(99%)={:.4}%, CVaR(99%)={:.4}% (from {} PnL samples)", + epoch + 1, self.hyperparams.epochs, + var_95 * 100.0, cvar_95 * 100.0, + var_99 * 100.0, cvar_99 * 100.0, + returns.len() + ); + } + + // Financial metrics + let epoch_sharpe = { + let financials = compute_epoch_financials( + &self.pnl_history, + &monitor.action_counts, + 100_000.0, + ); + training_metrics::set_epoch_financial_metrics( + "dqn", "current", + financials.sharpe, financials.sortino, financials.win_rate, + financials.max_drawdown, financials.profit_factor, + financials.total_return, financials.avg_return, + financials.total_trades as f64, + ); + training_metrics::set_epoch_action_distribution( + "dqn", "current", + financials.buy_pct, financials.sell_pct, financials.hold_pct, + ); + info!( + "Epoch {}/{}: Sharpe={:.2} WinRate={:.1}% MaxDD={:.3}% PF={:.2} Return={:+.2}% Trades={}", + epoch + 1, self.hyperparams.epochs, + financials.sharpe, financials.win_rate * 100.0, + financials.max_drawdown * 100.0, financials.profit_factor, + financials.total_return * 100.0, financials.total_trades, + ); + // QuestDB epoch row + { + let bps = if epoch_duration.as_secs_f64() > 0.0 { + train_step_count as f64 / epoch_duration.as_secs_f64() + } else { + 0.0 + }; + let buf_size = { + let agent = self.agent.read().await; + agent.get_replay_buffer_size().unwrap_or(0) as f64 + }; + let run_id = std::env::var("HOSTNAME").unwrap_or_else(|_| "local".to_owned()); + questdb_sink::record_training_epoch(&questdb_sink::EpochRecord { + model: "dqn", + fold: "current", + run_id: &run_id, + symbol: "", + epoch: (epoch + 1) as u32, + loss: avg_loss, + val_loss, + batches_per_second: bps, + epoch_duration_secs: epoch_duration.as_secs_f64(), + q_mean, + q_max, + gradient_norm: avg_grad_norm, + learning_rate: current_lr, + replay_buffer_size: buf_size, + action_entropy: epoch_entropy, + action_diversity: diversity_percentage / 100.0, + sharpe: financials.sharpe, + sortino: financials.sortino, + win_rate: financials.win_rate, + max_drawdown: financials.max_drawdown, + profit_factor: financials.profit_factor, + total_return: financials.total_return, + avg_return: financials.avg_return, + total_trades: financials.total_trades as f64, + }); + } + + financials.sharpe + }; + + // Track metrics for early stopping + self.loss_history.push(avg_loss); + self.q_value_history.push(avg_q_value); + self.val_loss_history.push(val_loss); + self.sharpe_history.push(epoch_sharpe); + + // Limit history vectors to prevent unbounded growth + const MAX_HISTORY_LEN: usize = 100; + if self.loss_history.len() > MAX_HISTORY_LEN { + self.loss_history.drain(0..50); + } + if self.q_value_history.len() > MAX_HISTORY_LEN { + self.q_value_history.drain(0..50); + } + if self.val_loss_history.len() > MAX_HISTORY_LEN { + self.val_loss_history.drain(0..50); + } + if self.sharpe_history.len() > MAX_HISTORY_LEN { + self.sharpe_history.drain(0..50); + } + + Ok(EpochLogOutput { + avg_loss, + avg_q_value, + avg_grad_norm, + epoch_sharpe, + val_loss, + q_min, + q_max, + q_mean, + }) + } + + // ═══════════════════════════════════════════════════════════════════════ + // Helper: Checkpoint saving + early stopping checks + // Returns Err for early stopping (caller propagates), Ok(()) to continue + // ═══════════════════════════════════════════════════════════════════════ + + pub(crate) async fn handle_epoch_checkpoints_and_early_stopping( + &mut self, + epoch: usize, + train_step_count: usize, + log_output: &EpochLogOutput, + checkpoint_callback: &mut F, + ) -> Result<()> + where + F: FnMut(usize, Vec, bool) -> Result + Send, + { + // C4 FIX: Save best model checkpoint when Sharpe improves + if train_step_count > 0 && log_output.epoch_sharpe > self.best_sharpe { + self.best_sharpe = log_output.epoch_sharpe; + self.best_val_loss = log_output.val_loss; + self.best_epoch = epoch + 1; + + info!( + "New best Sharpe: {:.4} at epoch {} (val_loss={:.6})", + log_output.epoch_sharpe, epoch + 1, log_output.val_loss, + ); + + let checkpoint_data = self.serialize_model().await?; + + // Verify checkpoint integrity + if self.safety_level != crate::safety::SafetyLevel::Permissive { + if checkpoint_data.is_empty() { + let msg = "SAFETY: Checkpoint verification failed - empty checkpoint data"; + match self.safety_level { + crate::safety::SafetyLevel::Strict => { + return Err(anyhow::anyhow!("{}", msg)); + }, + crate::safety::SafetyLevel::Normal => { + debug!("{} (continuing anyway)", msg); + }, + _ => {}, + } + } else { + debug!("SAFETY: Checkpoint verification passed ({} bytes)", checkpoint_data.len()); + } + } + + let ckpt_size = checkpoint_data.len() as f64; + let ckpt_start = std::time::Instant::now(); + let best_checkpoint_path = checkpoint_callback( + epoch + 1, checkpoint_data, true, + ).context("Failed to save best checkpoint")?; + training_metrics::record_checkpoint_save("dqn", "current", ckpt_start.elapsed().as_secs_f64(), ckpt_size); + + info!("Best model saved to: {}", best_checkpoint_path); + } + + // Early stopping checks (skip if no training occurred) + if train_step_count > 0 { + let old_should_stop = self.check_early_stopping(log_output.avg_q_value, epoch); + let patience_should_stop = if self.hyperparams.early_stopping_enabled + && epoch + 1 >= self.hyperparams.min_epochs_before_stopping { + self.early_stopping.should_stop(-log_output.epoch_sharpe) + } else { + false + }; + + if let Some(stop_reason) = old_should_stop { + warn!( + "Early stopping triggered at epoch {}/{}: {}", + epoch + 1, self.hyperparams.epochs, stop_reason + ); + info!( + "Final metrics: loss={:.6}, Q-value={:.4}", + log_output.avg_loss, log_output.avg_q_value + ); + + let checkpoint_data = self.serialize_model().await + .context("Failed to serialize model for early stopping checkpoint")?; + let checkpoint_size = checkpoint_data.len(); + let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false) + .context("Failed to save early stopping checkpoint")?; + info!( + "Early stopping checkpoint saved to: {} ({} bytes)", + checkpoint_path, checkpoint_size + ); + + return Err(anyhow::anyhow!( + "Training terminated by early stopping at epoch {}/{}: {}", + epoch + 1, self.hyperparams.epochs, stop_reason + )); + } + + if patience_should_stop { + warn!( + "WAVE 24 Early stopping (patience) triggered at epoch {}/{}! No improvement for {} epochs", + epoch + 1, self.hyperparams.epochs, self.hyperparams.gradient_collapse_patience + ); + info!("Best Sharpe: {:.4} at epoch {} (best val_loss proxy: {:.6})", + self.best_sharpe, self.best_epoch, self.early_stopping.get_best_val_loss()); + + let checkpoint_data = self.serialize_model().await?; + let checkpoint_path = checkpoint_callback(epoch + 1, checkpoint_data, false)?; + info!("Patience-based early stopping checkpoint saved to: {}", checkpoint_path); + + return Err(anyhow::anyhow!("Training terminated by patience-based early stopping at epoch {}", epoch + 1)); + } + } + + Ok(()) + } +}