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

feat(sp22): H6 Phase 3 α Phase C1 — collector W ptr setter

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-13 08:14:01 +02:00
parent 5d163e0e1d
commit 5106e3b117
4 changed files with 96 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<f32>,
pub(crate) w_aux_to_q_dir: cudarc::driver::CudaSlice<f32>,
/// 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<f32>,

59
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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)

16
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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

37
docs/dqn-wire-up-audit.md
View File

@@ -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.