From 5106e3b117e89ae5a7e6f0de9f09ed0fe600ad87 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 13 May 2026 08:14:01 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22):=20H6=20Phase=203=20=CE=B1=20Phase?= =?UTF-8?q?=20C1=20=E2=80=94=20collector=20W=20ptr=20setter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the trainer's `w_aux_to_q_dir [4]` device pointer into the GpuExperienceCollector so rollout-time action selection sees the active atom-shift for the direction branch. Changes: - gpu_experience_collector.rs: new field aux_w_to_q_dir_dev_ptr (u64, default 0) + setter set_aux_w_to_q_dir_ptr(). Both rollout launchers (compute_expected_q + quantile_q_select) now pass W ptr + exp_states_f32 + STATE_DIM_PADDED instead of NULL placeholders. NULL-safe via the kernels' existing aux_shift_active gating. - gpu_dqn_trainer.rs: w_aux_to_q_dir field promoted to pub(crate) for cross-module access via raw_ptr(). - trainers/dqn/trainer/training_loop.rs: new wire-up block after SP15 warm-count setter, mirroring the established setter pattern. NULL-safe on test scaffolds where fused_ctx or collector is absent. End-state — rollout activation: - compute_expected_q + quantile_q_select now return shifted E[Q] / quantile-blends per direction action during rollout. - With trainer's W trained each step (Step 8+11 Adam), the rollout policy's direction-action distribution actively reflects the learned aux→policy coupling. state_121's per-env value drives a per-(env, action) bias of magnitude W[a] (≤0.5 initial prior, learned thereafter). Verification: cargo check -p ml --lib clean (0 errors, 21 pre-existing warnings). Trainer + collector now smoke-ready end-to-end for the trainer/rollout side. Eval-side activation pending Phase D. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 2 +- .../cuda_pipeline/gpu_experience_collector.rs | 59 +++++++++++++------ .../src/trainers/dqn/trainer/training_loop.rs | 16 +++++ docs/dqn-wire-up-audit.md | 37 ++++++++++++ 4 files changed, 96 insertions(+), 18 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index babe9509c..15ffd2abf 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -6924,7 +6924,7 @@ pub struct GpuDqnTrainer { /// `[Short, Hold, Long, Flat]` — domain belief that aux conviction × /// position-sign biases toward aligned trades. Adam refines from /// the prior via projection backward in c51_grad_kernel. - w_aux_to_q_dir: cudarc::driver::CudaSlice, + pub(crate) w_aux_to_q_dir: cudarc::driver::CudaSlice, /// SP22 H6 Phase 3 α — Adam first moment (m) for `w_aux_to_q_dir`. /// Same shape [b0_size=4] f32; zero-init per Adam convention. adam_m_w_aux: cudarc::driver::CudaSlice, diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index aa7efacf5..45040f4bc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1390,6 +1390,14 @@ pub struct GpuExperienceCollector { /// Set via set_isv_signals_ptr(). 0 = NULL (static hold). isv_signals_dev_ptr: u64, + /// SP22 H6 Phase 3 α Phase C1 (2026-05-13): cached device pointer to + /// the trainer's `w_aux_to_q_dir [b0_size=4]` buffer. Default 0; + /// wired in via `set_aux_w_to_q_dir_ptr()` after trainer construction. + /// When non-zero, the rollout-time `compute_expected_q` + + /// `quantile_q_select` launchers thread W + `exp_states_f32` for + /// dir-branch atom-position shift `z_n + W[a]*state_121[b]`. + aux_w_to_q_dir_dev_ptr: u64, + /// SP15 Wave 2 (2026-05-06): trainer's `sp15_alpha_warm_count` /// mapped-pinned scratch buffer dev pointer ([1] f32). Set via /// `set_sp15_alpha_warm_count_ptr()` after trainer construction. 0 @@ -3211,6 +3219,7 @@ impl GpuExperienceCollector { cost_anneal_pinned, cost_anneal_dev_ptr, isv_signals_dev_ptr: 0, // NULL until trainer sets it + aux_w_to_q_dir_dev_ptr: 0, // SP22 H6 Phase 3 α Phase C1 — NULL until trainer sets it sp15_alpha_warm_count_dev_ptr: 0, // NULL until trainer wires it (SP15 Wave 2) // SP15 Phase 1.3.b-followup-B (2026-05-07): both per-env tiles // owned by this struct (collector), allocated above. The @@ -3303,6 +3312,19 @@ impl GpuExperienceCollector { self.isv_signals_dev_ptr = dev_ptr; } + /// SP22 H6 Phase 3 α Phase C1 (2026-05-13): set the trainer's + /// `w_aux_to_q_dir [4]` device pointer so the collector's rollout- + /// time `compute_expected_q` + `quantile_q_select` launchers apply + /// the active atom-shift `z_n_effective = z_n + W[a] * state_121[b]` + /// for direction-branch action selection. Pass 0 to disable + /// (test scaffold path; both kernels NULL-safe → bit-identical + /// pre-Phase-3-α behavior). The collector also reads its own + /// `exp_states_f32` rollout-state buffer for `state_121` so no + /// state-pointer setter is needed. + pub fn set_aux_w_to_q_dir_ptr(&mut self, dev_ptr: u64) { + self.aux_w_to_q_dir_dev_ptr = dev_ptr; + } + /// SP15 Wave 2 (2026-05-06): set the trainer's /// `sp15_alpha_warm_count` mapped-pinned scratch buffer dev_ptr /// ([1] f32). The collector calls this after construction (mirrors @@ -5646,16 +5668,18 @@ impl GpuExperienceCollector { unsafe { let null_atom_stats = 0u64; let null_atom_positions = 0u64; // NULL = use linear atom spacing - // SP22 H6 Phase 3 α atom-shift collector-side: Phase C1 will - // wire trainer's W ptr in via a setter (collector doesn't own - // the Adam-trained weight). For now NULL → kernel collapses - // shift to 0, bit-identical to pre-Phase-3-α rollout behavior. - let null_w_aux = 0u64; - let null_batch_states = 0u64; + // SP22 H6 Phase 3 α Phase C1 (2026-05-13): atom-shift + // active when trainer has wired its W ptr via + // set_aux_w_to_q_dir_ptr(). exp_states_f32 holds the + // current rollout-step's per-env states with stride + // STATE_DIM_PADDED (= 128). NULL-safe in-kernel when + // W ptr is 0 → shift collapses to 0 bit-identical. + let w_aux_dev_ptr = self.aux_w_to_q_dir_dev_ptr; + let exp_states_dev_ptr = self.exp_states_f32.raw_ptr(); let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; let state_dim_i32 = - ml_core::state_layout::STATE_DIM as i32; + ml_core::state_layout::STATE_DIM_PADDED as i32; self.stream .launch_builder(&self.expected_q_kernel) .arg(&self.exp_v_logits) @@ -5671,8 +5695,8 @@ impl GpuExperienceCollector { .arg(&null_atom_stats) .arg(&mut self.q_var_buf) .arg(&null_atom_positions) // atom_positions (NULL = linear) - .arg(&null_w_aux) - .arg(&null_batch_states) + .arg(&w_aux_dev_ptr) + .arg(&exp_states_dev_ptr) .arg(&aux_dir_prob_index) .arg(&state_dim_i32) .launch(launch_cfg) @@ -6040,13 +6064,14 @@ impl GpuExperienceCollector { let iqn_r = self.iqn_readiness; let util_e = self.util_ema; let support_dev_ptr = self.per_sample_support_buf.dev_ptr; - // SP22 H6 Phase 3 α atom-shift collector-side: Phase C1 will - // wire the trainer's W ptr via a setter. For now NULL → shift - // collapses to 0 → bit-identical to pre-Phase-3-α rollout. - let null_w_aux = 0u64; - let null_batch_states = 0u64; + // SP22 H6 Phase 3 α Phase C1 (2026-05-13): atom-shift + // active when trainer has wired its W ptr via + // set_aux_w_to_q_dir_ptr(). Shares the same exp_states_f32 + // rollout-state buffer as compute_expected_q above. + let w_aux_dev_ptr = self.aux_w_to_q_dir_dev_ptr; + let exp_states_dev_ptr = self.exp_states_f32.raw_ptr(); let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; - let state_dim_i32 = ml_core::state_layout::STATE_DIM as i32; + let state_dim_i32 = ml_core::state_layout::STATE_DIM_PADDED as i32; unsafe { self.stream .launch_builder(&self.quantile_q_select_kernel) @@ -6062,8 +6087,8 @@ impl GpuExperienceCollector { .arg(&support_dev_ptr) // [N, 4, 3] stride-12 per-sample per-branch support (mapped pinned) .arg(&iqn_r) .arg(&util_e) - .arg(&null_w_aux) - .arg(&null_batch_states) + .arg(&w_aux_dev_ptr) + .arg(&exp_states_dev_ptr) .arg(&aux_dir_prob_index) .arg(&state_dim_i32) .launch(launch_cfg) diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 365f87132..7c93ec3d4 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1139,6 +1139,22 @@ impl DQNTrainer { } } + // SP22 H6 Phase 3 α Phase C1 (2026-05-13): wire the trainer's + // `w_aux_to_q_dir [4]` device pointer into the collector so + // the rollout-time `compute_expected_q` + `quantile_q_select` + // launchers apply the active atom-shift for direction-branch + // action selection. Mirrors the `set_isv_signals_ptr` / + // `set_sp15_alpha_warm_count_ptr` patterns. NULL-safe path: + // when the trainer is absent (test scaffold), no wire-up + // happens → both kernels see the collector's default 0 → + // shift collapses to 0 (bit-identical pre-Phase-3-α). + if let Some(ref fused) = self.fused_ctx { + let w_aux_ptr = fused.trainer().w_aux_to_q_dir.raw_ptr(); + if let Some(ref mut collector) = self.gpu_experience_collector { + collector.set_aux_w_to_q_dir_ptr(w_aux_ptr); + } + } + // SP15 Wave 4.3 / Phase 3.5.5.b → Phase 1.3.b-followup-B // (2026-05-07): the `set_sp15_dd_trajectory_prev_dd_ptr` // wiring block was DELETED. The trainer's `[1]` scratch diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 4710f051c..dbf1ae8bb 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17177,3 +17177,40 @@ Both scratch buffers default to 0 → when `aux_shift_active` is false in c51_lo W now trains via the c51 loss gradient. Starting from structural prior `[-0.5, 0.0, +0.5, 0.0]`, Adam refines W per-action based on observed `state_121 × (loss reduction from atom-shift)`. The aux head trains independently via its supervised next-bar CE; together they form a learned cross-coupling from aux direction predictions to dir-branch Q distribution shifts. Next session can dispatch a smoke at this checkpoint to measure whether the adaptive W (vs. fixed structural prior) lifts WR off 50%. If WR moves on the trainer side without rollout-side activation, Phase C1 is the next priority. If not, the hypothesis needs reconsideration. + +#### Phase C1 — collector W ptr setter (2026-05-13) + +**Phase C1 — rollout-time atom-shift activation**: + +The collector's `compute_expected_q` + `quantile_q_select` launchers now read the trainer's `w_aux_to_q_dir [4]` device pointer (instead of NULL) when the trainer is wired in. + +**Changes**: + +- `gpu_experience_collector.rs`: + - New field `aux_w_to_q_dir_dev_ptr: u64` (default 0). + - New setter `set_aux_w_to_q_dir_ptr(dev_ptr: u64)` mirroring `set_isv_signals_ptr` + `set_sp15_alpha_warm_count_ptr` patterns. + - Both rollout-time launchers (`compute_expected_q` line ~5666, `quantile_q_select` line ~6041) now pass `self.aux_w_to_q_dir_dev_ptr` for W and `self.exp_states_f32.raw_ptr()` for batch_states, with `STATE_DIM_PADDED` (= 128) as the stride. + +- `gpu_dqn_trainer.rs`: + - `w_aux_to_q_dir` field promoted from private to `pub(crate)` so `training_loop.rs` can read its raw_ptr(). + +- `trainers/dqn/trainer/training_loop.rs`: + - New wire-up block right after the SP15 warm-count setter calls `collector.set_aux_w_to_q_dir_ptr(fused.trainer().w_aux_to_q_dir.raw_ptr())`. NULL-safe: skipped when fused_ctx or collector absent (test scaffolds). + +**End-state — rollout activation complete**: + +Action selection during rollout now sees atom-shifted Q values: +- `compute_expected_q` returns shifted E[Q] per direction action → policy / epsilon-greedy / Thompson sampling picks accordingly. +- `quantile_q_select` returns shifted quantile-blended Q per direction action → IQR-based action selection mixes the shift into its blend. + +With the trainer's W updated each step by Step 8+11 Adam, the rollout policy's direction-action distribution now actively reflects the learned aux→policy coupling. State_121's per-env value (recentered aux p_up) drives a per-(env, action) bias of magnitude W[a] (≤ 0.5 on initial prior, learned thereafter). + +**Verification**: +- `cargo check -p ml --lib`: 0 errors, 21 pre-existing warnings. + +**Remaining work**: +- **Phase D** (7 tasks): A2 eval-side aux infrastructure. Eval-time action selection doesn't see atom-shift yet (the eval backtester has its own forward + action-selection chain separate from the trainer/collector). +- **Phase E** (7 tasks): verification gates including capture-graph integrity, compute-sanitizer. +- **Phase F** (5 tasks): audit doc + atomic commit + push + smoke + verdict. + +**Trainer + collector are now smoke-ready end-to-end for the trainer/rollout side**. Eval-side activation (Phase D) is needed for the validation backtests to also benefit from atom-shift.