From 809295dc5f305f96391d30d0bf2f2875e1da7f69 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sun, 8 Mar 2026 15:11:01 +0100
Subject: [PATCH] =?UTF-8?q?fix(ml):=20correct=20GPU=20experience=20path=20?=
 =?UTF-8?q?=E2=80=94=20aligned=20state=5Fdim,=20reward=20tracking,=20Prome?=
 =?UTF-8?q?theus=20label?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three bugs in the GPU experience collection hot path:

1. gpu_batch_to_experiences() had hardcoded state_dim=43 but the CUDA kernel
   outputs states at the ALIGNED dimension (56 with OFI, 48 without). After
   sample 0, every replay buffer entry had corrupted state vectors — the
   network was learning from garbage data.

2. GPU path never called monitor.track_reward(), so mean_reward was always
   reported as 0.0 in epoch logs despite the agent generating real rewards.

3. Action tracking was double-counted (direct array write + track_action_by_exposure),
   inflating diversity metrics by 2x. Consolidated into single bounded call.

Also adds missing app.kubernetes.io/component label to job-template.yaml
so Prometheus training-pods scrape job discovers training pods.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/src/trainers/dqn/trainer.rs | 33 ++++++++++++++-------------
 infra/k8s/training/job-template.yaml  |  1 +
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs
index 076f45e3d..c8d9599fa 100644
--- a/crates/ml/src/trainers/dqn/trainer.rs
+++ b/crates/ml/src/trainers/dqn/trainer.rs
@@ -1910,25 +1910,26 @@ impl DQNTrainer {
 
                     match collector.collect_experiences(features_buf, targets_buf, &episode_starts, &config) {
                         Ok(batch) => {
-                            // Track GPU-collected actions in the epoch monitor so action
-                            // diversity / entropy metrics are computed correctly.
-                            // batch.actions contains i32 indices in 0..4 (5 exposure levels).
-                            for &action_idx in &batch.actions {
-                                let clamped = action_idx.clamp(0, 4) as usize;
-                                monitor.action_counts[clamped] += 1;
-                            }
-
-                            // Populate pnl_history from GPU-collected rewards for financial metrics.
-                            // Without this, compute_epoch_financials() sees an empty deque and
-                            // reports Trades=0, Sharpe=0.00 even when the agent is trading.
+                            // Feed GPU-collected rewards to both pnl_history (for Sharpe)
+                            // and monitor (for mean_reward / reward validation).
                             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);
+                            }
+                            // Track GPU-collected actions for diversity / entropy metrics.
+                            // batch.actions contains i32 indices in 0..4 (5 exposure levels).
+                            for &action_idx in &batch.actions {
+                                monitor.track_action_by_exposure(action_idx.clamp(0, 4) as usize);
                             }
 
-                            let experiences = gpu_batch_to_experiences(&batch);
+                            // GPU kernel outputs states at the ALIGNED dimension
+                            // (56 with OFI, 48 without) — must match kernel_dims.0.
+                            let raw_dim = if self.hyperparams.mbp10_data_dir.is_some() { 51 } else { 43 };
+                            let aligned_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_dim, &self.device);
+                            let experiences = gpu_batch_to_experiences(&batch, aligned_dim);
                             let count = experiences.len();
                             info!("GPU collected {} experiences ({} episodes × {} timesteps)",
                                 count, batch.n_episodes, batch.timesteps);
@@ -4476,14 +4477,14 @@ impl DQNTrainer {
 /// Convert a GPU `ExperienceBatch` (flat arrays from CUDA kernel) into `Vec<Experience>`
 /// for insertion into the DQN replay buffer.
 ///
-/// State dimension is 43 (40 market + 3 portfolio features, assembled in-kernel).
-/// Actions are mapped to `u8` for the `Experience` struct. Rewards are stored as
-/// fixed-point (×10000) by `Experience::new`.
+/// `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<Experience> {
-    let state_dim = 43;
     let total = batch.n_episodes * batch.timesteps;
     let mut experiences = Vec::with_capacity(total);
 
diff --git a/infra/k8s/training/job-template.yaml b/infra/k8s/training/job-template.yaml
index a4d04333b..0f6feda87 100644
--- a/infra/k8s/training/job-template.yaml
+++ b/infra/k8s/training/job-template.yaml
@@ -19,6 +19,7 @@ spec:
         gitlab.com/prometheus_path: "/metrics"
       labels:
         app.kubernetes.io/name: training
+        app.kubernetes.io/component: training-workflow
         app.kubernetes.io/part-of: foxhunt
         foxhunt/job-type: training
     spec: