feat(sp20): Phase 3 Task 3.4 — replay buffer schema for per-bar aux_conf
Lands the producer→ring→consumer schema plumbing for per-bar aux_conf
(the K=2 peak-softmax-above-uniform confidence value the Phase 5
Aux→Q gate consumes at the Bellman target). Atomic across all struct
boundaries per `feedback_no_partial_refactor`.
Data flow (TWO struct boundaries, not one — plan errata Gap 11):
ExperienceCollector → GpuExperienceBatch → insert_batch()
(.aux_conf field) (8th arg)
→ ring → sample() → GpuBatchPtrs → Trainer
(.aux_conf_ptr) (Phase 5 consumer)
Phase 5 lands the consumer (Aux→Q gate kernel at the Bellman target);
this commit is the producer-side schema plumbing only.
Spec §4.4 Phase 5 forward-reference contract:
At each replay batch step:
aux_conf = ptrs.aux_conf_ptr[k] # SAMPLED bar
threshold = ISV[AUX_CONF_THRESHOLD_INDEX] # [0.01, 0.20]
temp = ISV[AUX_GATE_TEMP_INDEX]
gate = sigmoid((aux_conf - threshold) / temp) # ∈ [0, 1]
Q_target = gate × Q_full + (1 - gate) × mean_a Q(s, a)
Producer (experience_env_step):
- 1 new kernel arg `aux_conf_per_sample: float* [N×L×cf_mult]`.
- Per-bar write at kernel entry: `aux_conf_per_sample[out_off] =
sp20_compute_per_bar_aux_conf_k2(aux_logits + i*K)`.
- CF slot write: `aux_conf_per_sample[cf_off] =
aux_conf_per_sample[out_off]` (CF state shares the same env state
so shares the same aux signal).
- NULL-tolerant: defaults to 0.0f (K=2 uniform-prior sentinel).
GpuExperienceCollector:
- +`aux_conf_per_sample: CudaSlice<f32>` field.
- alloc with `total_output * cf_mult` for both on-policy + CF.
- Threads as kernel arg of experience_env_step.
- Clones into `GpuExperienceBatch.aux_conf` at end of
collect_experiences_gpu (size `total = base_total × 2`).
GpuExperienceBatch: +`aux_conf: CudaSlice<f32>` field.
GpuReplayBuffer (crates/ml-dqn):
- +`GpuBatchPtrs.aux_conf_ptr: u64`.
- +ring storage `aux_conf: CudaSlice<f32>` [capacity].
- +sample buffer `sample_aux_conf: CudaSlice<f32>` [max_batch_size].
- +`trainer_aux_conf_ptr: u64` for direct path.
- +`set_trainer_aux_conf_ptr` setter + `trainer_aux_conf_ptr`
accessor (mirrors `set_isv_signals_ptr` pattern).
- `insert_batch` adds 8th arg `aux_conf_in: &CudaSlice<f32>` —
scatters via `scatter_insert_f32` (reuses the rewards/dones
scatter kernel).
- `sample_proportional` adds gather: direct-to-trainer if
`trainer_aux_conf_ptr != 0`, else fallback into
`sample_aux_conf`. Both paths populate `GpuBatchPtrs.aux_conf_ptr`.
Atomic caller updates (insert_batch arg from 7 to 8):
- training_loop.rs:2522 (production)
- gpu_residency.rs:77 (smoke)
- performance.rs:144 (smoke)
- training_stability.rs:154+201 (smoke ×2)
- gpu_per_integration_test.rs:128 (integration)
- sp15_phase1_oracle_tests.rs:2170+2189+2356 (oracle ×3)
- 3 inline tests in gpu_replay_buffer.rs
Tests:
- test_aux_conf_schema_round_trip (GPU): inserts 8 transitions with
distinct aux_conf [0.0, 0.05, ..., 0.35]; samples 1; asserts the
gathered slot value matches one of the inserted values
(round-trip integrity invariant).
- test_aux_conf_trainer_ptr_wiring_contract (CPU): asserts
set_trainer_aux_conf_ptr / trainer_aux_conf_ptr setter+accessor
behavior matches the SP15 Phase 3.5.5.b mirror pattern.
Verification:
SQLX_OFFLINE=true cargo check -p ml --tests --features cuda # green
SQLX_OFFLINE=true cargo check -p ml-dqn --tests --features cuda # green
SQLX_OFFLINE=true cargo test -p ml --lib sp20 # 21/21 pass
SQLX_OFFLINE=true cargo test -p ml-dqn test_aux_conf_trainer_ptr_wiring_contract # CPU pass
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -40,6 +40,17 @@ pub struct GpuBatchPtrs {
|
||||
/// `sample_aux_sign_labels` (not direct-to-trainer) until B1 lands the
|
||||
/// trainer destination buffer.
|
||||
pub aux_sign_labels_ptr: u64,
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf at the
|
||||
/// SAMPLED bar (f32). Range [0, 0.5] — K=2 peak-softmax-above-
|
||||
/// uniform confidence. 0.0 = no information (sentinel / aux head
|
||||
/// unwired). Producer: ring-buffer scatter from
|
||||
/// `GpuExperienceBatch.aux_conf` (per-bar value written by
|
||||
/// `experience_env_step`). Consumer: Phase 5 Aux→Q gate at the
|
||||
/// Bellman target — `gate = sigmoid((aux_conf -
|
||||
/// ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`. Gate ∈ [0, 1]
|
||||
/// blends `Q_target_full` (gate=1) with `Q_target_baseline =
|
||||
/// mean_a Q(s, a)` (gate=0) per spec §4.4.
|
||||
pub aux_conf_ptr: u64,
|
||||
pub batch_size: usize,
|
||||
pub state_dim: usize,
|
||||
}
|
||||
@@ -167,6 +178,16 @@ pub struct GpuReplayBuffer {
|
||||
/// Written by `scatter_insert_i32` in `insert_batch`, read by
|
||||
/// `gather_i32_scalar` in `sample_proportional`.
|
||||
aux_sign_labels: CudaSlice<i32>,
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf ring
|
||||
/// buffer `[capacity]` f32. Range [0, 0.5] — K=2 peak-softmax-
|
||||
/// above-uniform confidence. 0.0 = no information (sentinel).
|
||||
/// Written by `scatter_insert_f32` in `insert_batch` (mirrors the
|
||||
/// aux_sign_labels i32 path one slot above); read by
|
||||
/// `gather_f32_scalar` in `sample_proportional` →
|
||||
/// `GpuBatchPtrs.aux_conf_ptr`. Consumed by Phase 5 Aux→Q gate at
|
||||
/// the Bellman target (currently no consumer — Phase 5 lands the
|
||||
/// gate kernel; this commit is the producer-side schema plumbing).
|
||||
aux_conf: CudaSlice<f32>,
|
||||
write_cursor: usize, size: usize,
|
||||
max_priority: CudaSlice<f32>,
|
||||
pending_max_priority: Option<CudaSlice<f32>>,
|
||||
@@ -194,6 +215,13 @@ pub struct GpuReplayBuffer {
|
||||
/// `gather_i32_scalar` writes here from the ring; `GpuBatchPtrs.aux_sign_labels_ptr`
|
||||
/// points at this buffer until B1 wires a trainer-resident destination.
|
||||
sample_aux_sign_labels: CudaSlice<i32>,
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-batch aux_conf gather
|
||||
/// destination. `gather_f32_scalar` writes here from the ring;
|
||||
/// `GpuBatchPtrs.aux_conf_ptr` points at this buffer until Phase 5
|
||||
/// wires a trainer-resident destination (the trainer's
|
||||
/// `aux_conf_at_state_buf` mirroring the
|
||||
/// `trainer_aux_sign_labels_ptr` direct-path pattern).
|
||||
sample_aux_conf: CudaSlice<f32>,
|
||||
total_sum_buf: CudaSlice<f32>,
|
||||
// Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step)
|
||||
update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator
|
||||
@@ -240,6 +268,17 @@ pub struct GpuReplayBuffer {
|
||||
/// actions / rewards. Fallback gather into `sample_aux_sign_labels` is
|
||||
/// skipped on the direct path so we don't double-write.
|
||||
trainer_aux_sign_labels_ptr: u64,
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): trainer destination for the
|
||||
/// aux_conf f32 column. When non-zero, `sample_proportional`
|
||||
/// gathers `aux_conf_at_state` directly into the trainer's
|
||||
/// (Phase-5-allocated) `aux_conf_at_state_buf`. When 0 (default
|
||||
/// at construction), `sample_proportional` falls back to gathering
|
||||
/// into `sample_aux_conf` and `GpuBatchPtrs.aux_conf_ptr` points
|
||||
/// there. Same direct-path pattern as `trainer_aux_sign_labels_ptr`
|
||||
/// — Phase 5 wires the trainer-side destination via a
|
||||
/// `set_trainer_aux_conf_ptr` call (mirrors the existing
|
||||
/// `set_trainer_buffers` 8-arg setter).
|
||||
trainer_aux_conf_ptr: u64,
|
||||
trainer_state_dim_padded: usize,
|
||||
/// C1/P1: Cached learning_health value from the trainer (updated once per epoch).
|
||||
/// Controls whether priority updates use the standard or diversity-weighted kernel.
|
||||
@@ -361,6 +400,11 @@ impl GpuReplayBuffer {
|
||||
// Zero-init in B0 — B1 producer kernel writes -1/0/1.
|
||||
let aux_lbl = a32i(stream, cap, "aux_sign_labels")?;
|
||||
let s_aux = a32i(stream, mbs, "s_aux_sign_labels")?;
|
||||
// SP20 Phase 3 Task 3.4 (2026-05-10): aux_conf ring + per-batch
|
||||
// gather destination. Zero-init = K=2 uniform-prior sentinel
|
||||
// (no information). Mirrors the aux_sign_labels f32 path.
|
||||
let aux_conf_ring = a32f(stream, cap, "aux_conf")?;
|
||||
let s_aux_conf = a32f(stream, mbs, "s_aux_conf")?;
|
||||
|
||||
// Pre-allocate PER update scratch buffers (zero cuMemAlloc per step)
|
||||
let update_batch_max = a32f(stream, 1, "update_batch_max")?;
|
||||
@@ -412,6 +456,7 @@ impl GpuReplayBuffer {
|
||||
states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
|
||||
episode_ids: ep_ids,
|
||||
aux_sign_labels: aux_lbl,
|
||||
aux_conf: aux_conf_ring,
|
||||
write_cursor: 0, size: 0, max_priority: mp,
|
||||
pending_max_priority: None, current_step: 0,
|
||||
priorities_pa, prefix_sum,
|
||||
@@ -423,6 +468,7 @@ impl GpuReplayBuffer {
|
||||
sample_priorities: sp, sample_weights: sw,
|
||||
sample_episode_ids: s_ep,
|
||||
sample_aux_sign_labels: s_aux,
|
||||
sample_aux_conf: s_aux_conf,
|
||||
sample_max_weight: smw, total_sum_buf: tsb,
|
||||
update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare),
|
||||
_mean_action_pinned: mean_action_pinned, mean_action_dev_ptr,
|
||||
@@ -441,6 +487,7 @@ impl GpuReplayBuffer {
|
||||
trainer_dones_ptr: 0,
|
||||
trainer_is_weights_ptr: 0,
|
||||
trainer_aux_sign_labels_ptr: 0,
|
||||
trainer_aux_conf_ptr: 0,
|
||||
trainer_state_dim_padded: 0,
|
||||
learning_health_cache: 1.0,
|
||||
// SP15 Phase 3.5.5.b: NULL until trainer wires the ISV bus
|
||||
@@ -611,6 +658,7 @@ impl GpuReplayBuffer {
|
||||
/// scatter-insert as f32 directly.
|
||||
/// #30: Accept f32 states from the experience collector.
|
||||
/// States are stored as f32 in the replay buffer (no f32 truncation).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn insert_batch(
|
||||
&mut self,
|
||||
sf: &CudaSlice<f32>, // f32 states from experience collector
|
||||
@@ -619,6 +667,7 @@ impl GpuReplayBuffer {
|
||||
rw: &CudaSlice<f32>,
|
||||
dn: &CudaSlice<f32>,
|
||||
aux_sign: &CudaSlice<i32>, // SP13 Layer B: 30-bar aux sign label, -1/0/1
|
||||
aux_conf_in: &CudaSlice<f32>, // SP20 Phase 3 Task 3.4: per-bar aux_conf, [0, 0.5]
|
||||
bs: usize,
|
||||
) -> Result<(), MLError> {
|
||||
if bs == 0 { return Ok(()); }
|
||||
@@ -667,6 +716,16 @@ impl GpuReplayBuffer {
|
||||
.arg(&self.aux_sign_labels).arg(&s_aux_in).arg(&ci).arg(&cpi).arg(&bsi)
|
||||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux i32: {e}")))?;
|
||||
}
|
||||
// SP20 Phase 3 Task 3.4 (2026-05-10): aux_conf scatter
|
||||
// (f32, [0, 0.5]). Mirrors the aux_sign_labels i32 path
|
||||
// immediately above; reuses `scatter_insert_f32` (the same
|
||||
// kernel that scatters rewards / dones).
|
||||
let s_aux_conf_in = if off > 0 { aux_conf_in.slice(off..) } else { aux_conf_in.slice(0..) };
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
|
||||
.arg(&self.aux_conf).arg(&s_aux_conf_in).arg(&ci).arg(&cpi).arg(&bsi)
|
||||
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux_conf f32: {e}")))?;
|
||||
}
|
||||
}
|
||||
// Broadcast max_priority into pre-allocated prio buffer (raw
|
||||
// base priority, fed as `raw_priorities` to per_insert_pa).
|
||||
@@ -962,6 +1021,39 @@ impl GpuReplayBuffer {
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3d (SP20 Phase 3 Task 3.4): gather per-bar aux_conf (f32).
|
||||
// Always gathers into `sample_aux_conf` (the GpuBatchPtrs
|
||||
// fallback path) UNLESS the trainer wired
|
||||
// `trainer_aux_conf_ptr` via `set_trainer_aux_conf_ptr` (Phase
|
||||
// 5 lands the trainer destination buffer). When wired, gathers
|
||||
// directly into the trainer's destination — same direct-path
|
||||
// pattern as `aux_sign_labels` above.
|
||||
let aux_conf_dst_ptr = if self.trainer_aux_conf_ptr != 0 {
|
||||
self.trainer_aux_conf_ptr
|
||||
} else {
|
||||
self.sample_aux_conf.raw_ptr()
|
||||
};
|
||||
if self.trainer_aux_conf_ptr != 0 {
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.kernels.gather_f32_scalar)
|
||||
.arg(&aux_conf_dst_ptr)
|
||||
.arg(&self.aux_conf)
|
||||
.arg(&self.sample_indices_i64)
|
||||
.arg(&bsi)
|
||||
.launch(lcfg(batch_size))
|
||||
.map_err(|e| MLError::ModelError(format!("g aux_conf direct: {e}")))?;
|
||||
}
|
||||
} else {
|
||||
unsafe {
|
||||
let cap_i32 = self.capacity() as i32;
|
||||
self.stream.launch_builder(&self.kernels.gather_f32)
|
||||
.arg(&mut self.sample_aux_conf).arg(&self.aux_conf)
|
||||
.arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
|
||||
.launch(lcfg(batch_size))
|
||||
.map_err(|e| MLError::ModelError(format!("g aux_conf fallback: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Step 6: IS weights via GPU-resident total_sum (zero CPU readback)
|
||||
// is_weights_f32 reads total_sum from GPU pointer.
|
||||
// When direct_to_trainer is active, write IS weights to trainer's is_weights_buf.
|
||||
@@ -1035,6 +1127,12 @@ impl GpuReplayBuffer {
|
||||
// (already populated by the direct-path gather above); the
|
||||
// fallback `sample_aux_sign_labels` is unused on this path.
|
||||
aux_sign_labels_ptr: self.trainer_aux_sign_labels_ptr,
|
||||
// SP20 Phase 3 Task 3.4: aux_conf points at the trainer
|
||||
// destination if wired (Phase 5), otherwise at the
|
||||
// fallback `sample_aux_conf` (already populated by the
|
||||
// gather block above — both paths populate, the only
|
||||
// difference is the destination pointer).
|
||||
aux_conf_ptr: aux_conf_dst_ptr,
|
||||
batch_size,
|
||||
state_dim: sd,
|
||||
})
|
||||
@@ -1049,12 +1147,32 @@ impl GpuReplayBuffer {
|
||||
indices_ptr: self.sample_indices_u32.raw_ptr(),
|
||||
episode_ids_ptr: self.sample_episode_ids.raw_ptr(),
|
||||
aux_sign_labels_ptr: self.sample_aux_sign_labels.raw_ptr(),
|
||||
aux_conf_ptr: aux_conf_dst_ptr,
|
||||
batch_size,
|
||||
state_dim: sd,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): wire the trainer's
|
||||
/// destination pointer for aux_conf gathering. When non-zero,
|
||||
/// `sample_proportional` gathers `aux_conf_at_state` directly into
|
||||
/// the trainer's buffer (mirrors `set_trainer_buffers`'s
|
||||
/// `aux_sign_labels_ptr` direct-path semantic). When 0 (default
|
||||
/// at construction), falls back to the intermediate
|
||||
/// `sample_aux_conf` buffer. Phase 5 will call this once after
|
||||
/// the trainer's `aux_conf_at_state_buf` is allocated.
|
||||
pub fn set_trainer_aux_conf_ptr(&mut self, ptr: u64) {
|
||||
self.trainer_aux_conf_ptr = ptr;
|
||||
}
|
||||
|
||||
/// SP20 Phase 3 Task 3.4 accessor: returns the wired trainer
|
||||
/// aux_conf destination pointer (0 if unwired). Used by oracle
|
||||
/// tests to verify the wiring contract.
|
||||
pub const fn trainer_aux_conf_ptr(&self) -> u64 {
|
||||
self.trainer_aux_conf_ptr
|
||||
}
|
||||
|
||||
/// Reference to pre-allocated sample indices buffer (valid after sample_proportional).
|
||||
///
|
||||
/// Returns a view sized to the most recent `sample_proportional` batch size.
|
||||
@@ -1404,10 +1522,11 @@ mod tests {
|
||||
let mut r = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let d = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
// Non-trivial rewards drive non-zero priorities through `per_insert_pa`.
|
||||
let rewards_host: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0) * 0.05).collect();
|
||||
stream.memcpy_htod(&rewards_host, &mut r).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, n).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, n).unwrap();
|
||||
assert_eq!(b.len(), n, "post-insert size should be n");
|
||||
assert!(!b.is_empty(), "post-insert buffer should not be empty");
|
||||
assert!(b.can_sample(bs), "post-insert sampler should be ready");
|
||||
@@ -1483,10 +1602,11 @@ mod tests {
|
||||
let mut r = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let d = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
// Give non-zero rewards so priorities are non-zero after update
|
||||
let rewards_host: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0) * 0.01).collect();
|
||||
stream.memcpy_htod(&rewards_host, &mut r).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, n).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, n).unwrap();
|
||||
assert_eq!(b.len(), n);
|
||||
|
||||
// sample_proportional triggers per_prefix_scan — must complete, not hang
|
||||
@@ -1549,7 +1669,8 @@ mod tests {
|
||||
let r = stream.alloc_zeros::<f32>(bs).unwrap();
|
||||
let d = stream.alloc_zeros::<f32>(bs).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(bs).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, bs).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(bs).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, bs).unwrap();
|
||||
|
||||
// 2. Fire `update_priorities_gpu` with TD errors spanning [td_min, td_max].
|
||||
// Health stays at default 1.0, so the standard `per_update_pa` path
|
||||
@@ -1628,4 +1749,126 @@ mod tests {
|
||||
this assertion fires near the pre-fix value, the bug is back"
|
||||
);
|
||||
}
|
||||
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10) — replay-buffer schema test
|
||||
/// for the per-bar `aux_conf` field (Phase 5 Aux→Q gate input).
|
||||
///
|
||||
/// Asserts the producer→ring→consumer round-trip:
|
||||
/// 1. `insert_batch` accepts `aux_conf_in` arg + scatters into
|
||||
/// the ring's `aux_conf` column.
|
||||
/// 2. `sample_proportional` returns `GpuBatchPtrs.aux_conf_ptr`
|
||||
/// pointing at a valid f32 buffer.
|
||||
/// 3. The fallback path (no trainer dest wired) populates
|
||||
/// `sample_aux_conf` and the ptr matches its `raw_ptr()`.
|
||||
///
|
||||
/// Behavioral round-trip: insert N transitions with known aux_conf
|
||||
/// values; sample 1; the gathered slot must contain ONE of the
|
||||
/// inserted values (binomial sampling — not necessarily the first
|
||||
/// slot, but always SOMETHING from the set we inserted).
|
||||
///
|
||||
/// Per `pearl_tests_must_prove_not_lock_observations`: the
|
||||
/// assertion is "gathered value matches at least one inserted
|
||||
/// value" (round-trip integrity invariant), not "gathered value =
|
||||
/// some specific observed-from-a-run value".
|
||||
#[test]
|
||||
fn test_aux_conf_schema_round_trip() {
|
||||
let stream = make_stream();
|
||||
let cfg = GpuReplayBufferConfig {
|
||||
capacity: 64, alpha: 0.6, beta_start: 0.4, beta_max: 1.0,
|
||||
beta_annealing_steps: 100, epsilon: 1e-6,
|
||||
max_memory_bytes: 4 << 30, max_batch_size: 16,
|
||||
};
|
||||
let mut buf = GpuReplayBuffer::new(cfg, &stream).expect("buf");
|
||||
let sd = ml_core::state_layout::STATE_DIM;
|
||||
let n = 8;
|
||||
|
||||
let s = stream.alloc_zeros::<f32>(n * sd).unwrap();
|
||||
let ns = stream.alloc_zeros::<f32>(n * sd).unwrap();
|
||||
let a = stream.alloc_zeros::<u32>(n).unwrap();
|
||||
let mut r = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let d = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let aux_sign = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
|
||||
// Non-trivial rewards drive non-zero priorities.
|
||||
let rewards_host: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0) * 0.1).collect();
|
||||
stream.memcpy_htod(&rewards_host, &mut r).unwrap();
|
||||
|
||||
// Engineering aux_conf values: use distinct in-range values so
|
||||
// the round-trip can distinguish them. Range [0, 0.5] from the
|
||||
// K=2 peak-softmax-above-uniform definition.
|
||||
let aux_conf_host: Vec<f32> = (0..n)
|
||||
.map(|i| (i as f32) * 0.05) // 0.0, 0.05, 0.10, ..., 0.35
|
||||
.collect();
|
||||
let mut aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap();
|
||||
|
||||
buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, n)
|
||||
.expect("insert_batch with aux_conf");
|
||||
assert_eq!(buf.len(), n, "post-insert size = n");
|
||||
|
||||
// Sample 1 — fallback path (no trainer dest wired).
|
||||
let bs = 1;
|
||||
let ptrs = buf.sample_proportional(bs).expect("sample 1");
|
||||
|
||||
// Schema assertion 1: GpuBatchPtrs.aux_conf_ptr is non-zero
|
||||
// (points at a valid f32 buffer).
|
||||
assert!(ptrs.aux_conf_ptr != 0,
|
||||
"GpuBatchPtrs.aux_conf_ptr must be non-zero post-sample");
|
||||
|
||||
// Schema assertion 2: trainer destination unwired ⇒ ptr equals
|
||||
// sample_aux_conf.raw_ptr() (the fallback path).
|
||||
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
|
||||
"trainer_aux_conf_ptr should be 0 (unwired in this test)");
|
||||
|
||||
// Behavioral assertion: read back the gathered aux_conf and
|
||||
// verify it matches one of the inserted values (round-trip
|
||||
// integrity invariant).
|
||||
let mut gathered_host = vec![0.0_f32; bs];
|
||||
let gathered_view = buf.sample_aux_conf.slice(..bs);
|
||||
stream.memcpy_dtoh(&gathered_view, &mut gathered_host)
|
||||
.expect("dtoh sample_aux_conf");
|
||||
let gathered = gathered_host[0];
|
||||
let in_set = aux_conf_host.iter().any(|&v| (v - gathered).abs() < 1e-6);
|
||||
assert!(in_set,
|
||||
"Gathered aux_conf {gathered} not in inserted set {aux_conf_host:?} \
|
||||
— round-trip integrity violation");
|
||||
}
|
||||
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10) — trainer-aux_conf-ptr
|
||||
/// wiring contract test.
|
||||
///
|
||||
/// Asserts the `set_trainer_aux_conf_ptr` / `trainer_aux_conf_ptr`
|
||||
/// accessor pair behaves like `set_isv_signals_ptr` /
|
||||
/// `isv_signals_dev_ptr` (mirror of the SP15 Phase 3.5.5.b
|
||||
/// pattern). When wired, `sample_proportional` returns
|
||||
/// `aux_conf_ptr = trainer_aux_conf_ptr`, NOT the fallback
|
||||
/// `sample_aux_conf.raw_ptr()`.
|
||||
#[test]
|
||||
fn test_aux_conf_trainer_ptr_wiring_contract() {
|
||||
let stream = make_stream();
|
||||
let cfg = GpuReplayBufferConfig {
|
||||
capacity: 32, alpha: 0.6, beta_start: 0.4, beta_max: 1.0,
|
||||
beta_annealing_steps: 100, epsilon: 1e-6,
|
||||
max_memory_bytes: 4 << 30, max_batch_size: 8,
|
||||
};
|
||||
let mut buf = GpuReplayBuffer::new(cfg, &stream).expect("buf");
|
||||
|
||||
// Default: unwired.
|
||||
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
|
||||
"trainer_aux_conf_ptr defaults to 0 (unwired)");
|
||||
|
||||
// Wire to a fake non-zero pointer (any non-zero value passes
|
||||
// the wiring contract — the actual ptr would be the trainer's
|
||||
// `aux_conf_at_state_buf.raw_ptr()` in production wiring).
|
||||
let fake_ptr: u64 = 0x12345678;
|
||||
buf.set_trainer_aux_conf_ptr(fake_ptr);
|
||||
assert_eq!(buf.trainer_aux_conf_ptr(), fake_ptr,
|
||||
"trainer_aux_conf_ptr should reflect the set value");
|
||||
|
||||
// Re-zero (un-wire) — defensive, ensures the setter is not
|
||||
// sticky.
|
||||
buf.set_trainer_aux_conf_ptr(0);
|
||||
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
|
||||
"trainer_aux_conf_ptr should re-zero on set(0)");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2249,7 +2249,30 @@ extern "C" __global__ void experience_env_step(
|
||||
int hold_buffer_size,
|
||||
int aux_k,
|
||||
int hold_cost_scale_idx,
|
||||
int hold_reward_ema_idx
|
||||
int hold_reward_ema_idx,
|
||||
/* SP20 Phase 3 Task 3.4 (2026-05-10) — replay-buffer schema
|
||||
* `aux_conf` per-bar field for the Phase 5 Aux→Q gate.
|
||||
*
|
||||
* Layout: `[N × L × cf_mult]` f32 row-major. The Phase 5 gate at
|
||||
* the Bellman-target consumer needs `aux_conf_at_state` (the
|
||||
* confidence at the SAMPLED bar, not the current-step bar), so
|
||||
* this column threads through the replay-buffer scatter/gather
|
||||
* pipeline alongside `aux_sign_labels`.
|
||||
*
|
||||
* Producer: THIS kernel writes
|
||||
* `aux_conf_per_sample[out_off] = sp20_compute_per_bar_aux_conf_k2(
|
||||
* aux_logits_per_env + i * aux_k)` at the end of the kernel
|
||||
* body (after the reward composition, before the FINAL `out_rewards`
|
||||
* write at line ~3674). NULL-tolerant: when
|
||||
* `aux_logits_per_env == NULL` (test scaffolds), the slot retains
|
||||
* the kernel-entry default `0.0f` (matches the K=2 uniform-prior
|
||||
* `aux_conf = 0.5 - 0.5 = 0` semantic — no information).
|
||||
*
|
||||
* Consumer: replay buffer scatter (gpu_replay_buffer::insert_batch),
|
||||
* gather (sample_proportional → GpuBatchPtrs.aux_conf_ptr), and
|
||||
* eventual Phase 5 Aux→Q gate read in the Bellman-target loss
|
||||
* computation. */
|
||||
float* __restrict__ aux_conf_per_sample
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= N) return;
|
||||
@@ -2318,6 +2341,34 @@ extern "C" __global__ void experience_env_step(
|
||||
traded_per_sample[out_off] = 0;
|
||||
hold_at_exit_per_sample[out_off] = 0.0f;
|
||||
|
||||
/* SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf write for
|
||||
* the replay-buffer schema (Phase 5 Aux→Q gate consumer). Computed
|
||||
* ONCE per (i, t) at the K=2 uniform-prior or aux-head value;
|
||||
* value matches the per-bar branches' aux_conf_i computation in
|
||||
* Task 3.2 dual emission. Written unconditionally so every (i, t)
|
||||
* slot has a deterministic value (including segment_complete bars
|
||||
* and slots reached by early-return paths).
|
||||
*
|
||||
* Default 0.0 = K=2 uniform prior `peak_softmax = 0.5 ⇒ aux_conf =
|
||||
* 0.0` (no information). Real value computed when
|
||||
* `aux_logits_per_env != NULL` AND `aux_conf_per_sample != NULL`
|
||||
* (test-scaffold NULL safety). The Phase 5 gate read site at the
|
||||
* Bellman target uses
|
||||
* `gate = sigmoid((aux_conf - threshold) / temp)` — at default
|
||||
* 0.0 with threshold ∈ [0.01, 0.20] gate ≈ 0 ⇒ Q_target falls
|
||||
* back to `mean_a Q(s, a)` (the spec §4.4 "uncertain-state
|
||||
* neutralizer" baseline). */
|
||||
if (aux_conf_per_sample != NULL) {
|
||||
if (aux_logits_per_env != NULL) {
|
||||
const float* logits_row_init = aux_logits_per_env
|
||||
+ (long long)i * (long long)aux_k;
|
||||
aux_conf_per_sample[out_off] =
|
||||
sp20_compute_per_bar_aux_conf_k2(logits_row_init);
|
||||
} else {
|
||||
aux_conf_per_sample[out_off] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
/* Task 0.8 reward-contribution defaults — written BEFORE any early-return
|
||||
* so every (i,t) slot owned by this thread gets a deterministic value.
|
||||
* The cf_flip slot lives at cf_off (out_off + N*L); default it here too.
|
||||
@@ -4438,6 +4489,15 @@ extern "C" __global__ void experience_env_step(
|
||||
out_actions[cf_off] = cf_action;
|
||||
out_rewards[cf_off] = cf_reward_weighted;
|
||||
out_dones[cf_off] = (float)done;
|
||||
/* SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf at the
|
||||
* counterfactual slot. The CF state at cf_off shares the SAME
|
||||
* env state (only the action differs), so it shares the SAME
|
||||
* aux signal — mirror the on-policy write at out_off (above
|
||||
* at the kernel-entry default block). NULL-tolerant: same
|
||||
* guard as the out_off write site. */
|
||||
if (aux_conf_per_sample != NULL) {
|
||||
aux_conf_per_sample[cf_off] = aux_conf_per_sample[out_off];
|
||||
}
|
||||
/* C.2: component [1] = counterfactual reward magnitude for this sample.
|
||||
* Written at out_off (on-policy slot), not cf_off, so the EMA kernel's
|
||||
* [N*L, 6] reduction over out_off slots captures it alongside the other
|
||||
|
||||
@@ -477,6 +477,26 @@ pub struct GpuExperienceBatch {
|
||||
/// buffer threads this column through `insert_batch` → `gather_i32_scalar`
|
||||
/// → `GpuBatchPtrs.aux_sign_labels_ptr`.
|
||||
pub aux_sign_labels: CudaSlice<i32>,
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf per
|
||||
/// transition `[total]` on GPU (f32) — the K=2 peak-softmax-above-
|
||||
/// uniform `aux_conf ∈ [0, 0.5]` value at the time the bar was
|
||||
/// rolled out. Written by `experience_env_step` per (i, t) using
|
||||
/// `sp20_compute_per_bar_aux_conf_k2(aux_logits_per_env + i * 2)`.
|
||||
/// Default 0.0 = K=2 uniform-prior (no information) for test
|
||||
/// scaffolds where `aux_logits_per_env == NULL`.
|
||||
///
|
||||
/// The replay buffer threads this column through `insert_batch` →
|
||||
/// `gather_f32_scalar` → `GpuBatchPtrs.aux_conf_ptr` so the Phase
|
||||
/// 5 Aux→Q gate can read `aux_conf_at_state` (the confidence at
|
||||
/// the SAMPLED bar) inside the Bellman target computation:
|
||||
///
|
||||
/// gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])
|
||||
/// Q_target = gate × Q_full + (1 - gate) × Q_baseline
|
||||
///
|
||||
/// Per spec §4.4 the gate ∈ [0, 1] gracefully degrades the Bellman
|
||||
/// target toward `mean_a Q(s, a)` at uncertain states without
|
||||
/// freezing learning.
|
||||
pub aux_conf: CudaSlice<f32>,
|
||||
/// State dimensionality (so the consumer can reshape)
|
||||
pub state_dim: usize,
|
||||
/// Number of episodes
|
||||
@@ -849,6 +869,23 @@ pub struct GpuExperienceCollector {
|
||||
step_ret_per_sample: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps]
|
||||
trade_close_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
|
||||
trade_profitable_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
|
||||
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf scratch for
|
||||
/// the replay-buffer schema. Layout `[alloc_episodes ×
|
||||
/// alloc_timesteps × cf_mult = 2]` f32 row-major. Producer:
|
||||
/// `experience_env_step` writes
|
||||
/// `sp20_compute_per_bar_aux_conf_k2(aux_logits_per_env + i * 2)`
|
||||
/// at every (i, t) slot (mirrors the H6/Task-0.8 per-sample
|
||||
/// race-free pattern — each thread writes only its own out_off
|
||||
/// slot). Default 0.0 = K=2 uniform-prior. Consumer: assembled
|
||||
/// into `GpuExperienceBatch.aux_conf` at the end of
|
||||
/// `collect_experiences_gpu` (post-counterfactual) and threaded
|
||||
/// through `insert_batch` → ring → `gather_f32_scalar` →
|
||||
/// `GpuBatchPtrs.aux_conf_ptr`. Allocation size matches the other
|
||||
/// per-sample buffers (`alloc_episodes * alloc_timesteps`); the
|
||||
/// counterfactual extension `× cf_mult` happens at the assembly
|
||||
/// site (mirrors `aux_sign_labels` which is also `[total]` =
|
||||
/// `alloc_episodes * alloc_timesteps * 2`).
|
||||
aux_conf_per_sample: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps]
|
||||
/// Task 0.3 — per-sample distributional variance Kelly scale written by
|
||||
/// `experience_env_step` at line 1422 (`1 / (1 + sqrt(Var[Q]))`). Kernel
|
||||
/// writes only when `q_variance` is non-NULL AND `total_actions_for_var > 0`,
|
||||
@@ -1972,6 +2009,23 @@ impl GpuExperienceCollector {
|
||||
let trade_profitable_per_sample = stream
|
||||
.alloc_zeros::<i32>(total_output)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc trade_profitable_per_sample: {e}")))?;
|
||||
// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf scratch.
|
||||
// Sized `[total_output × cf_mult]` so the kernel writes both
|
||||
// on-policy (out_off) AND counterfactual (cf_off) slots — the
|
||||
// CF state at cf_off shares the SAME aux signal (same env
|
||||
// state, just different action), so both slots get the same
|
||||
// value, mirroring `cf_flip_per_sample`'s `total_output ×
|
||||
// cf_mult` sizing pattern (`gpu_experience_collector.rs` line
|
||||
// ~1945). alloc_zeros → uninitialised slots read as 0.0f =
|
||||
// K=2 uniform-prior (no information). Kernel overwrites at
|
||||
// every reached (i, t) with `sp20_compute_per_bar_aux_conf_k2(
|
||||
// aux_logits + i*K)`; un-reached slots stay 0.0f. Used
|
||||
// downstream by the Phase 5 Aux→Q gate at `gate =
|
||||
// sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) /
|
||||
// ISV[AUX_GATE_TEMP])`.
|
||||
let aux_conf_per_sample = stream
|
||||
.alloc_zeros::<f32>(total_output * cf_mult)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc aux_conf_per_sample: {e}")))?;
|
||||
// Task 0.3: per-sample Kelly-from-atoms variance scale, written by
|
||||
// experience_env_step (~line 1422). Alloc_zeros → uninitialised slots
|
||||
// read as 0.0f, which the host reducer filters out (only positive
|
||||
@@ -2900,6 +2954,7 @@ impl GpuExperienceCollector {
|
||||
action_mag_per_sample,
|
||||
step_ret_per_sample,
|
||||
trade_close_per_sample,
|
||||
aux_conf_per_sample,
|
||||
trade_profitable_per_sample,
|
||||
var_scale_per_sample,
|
||||
slot_live_per_sample,
|
||||
@@ -4899,6 +4954,20 @@ impl GpuExperienceCollector {
|
||||
|
||||
self.last_experience_count = total;
|
||||
|
||||
// SP20 Phase 3 Task 3.4 (2026-05-10): clone the per-bar
|
||||
// aux_conf scratch into a fresh `[total]` CudaSlice for the
|
||||
// batch struct. The kernel writes at `out_off` + `cf_off =
|
||||
// out_off + N*L` so the first `total = N*L*2` slots of the
|
||||
// scratch contain on-policy + CF; same dtod_clone pattern as
|
||||
// states/rewards/actions/dones above. NULL-tolerant
|
||||
// downstream: if aux_logits_per_env was NULL during the
|
||||
// rollout, the kernel left these slots at 0.0f (K=2 uniform
|
||||
// prior, no information) — Phase 5 gate will fall back to
|
||||
// `mean_a Q(s, a)` baseline at sample time.
|
||||
let aux_conf = dtod_clone_f32_native(
|
||||
&self.stream, &self.aux_conf_per_sample, total, "aux_conf"
|
||||
)?;
|
||||
|
||||
Ok(GpuExperienceBatch {
|
||||
states,
|
||||
next_states,
|
||||
@@ -4907,6 +4976,7 @@ impl GpuExperienceCollector {
|
||||
dones,
|
||||
episode_ids,
|
||||
aux_sign_labels,
|
||||
aux_conf,
|
||||
state_dim: sd,
|
||||
n_episodes,
|
||||
timesteps,
|
||||
@@ -6331,6 +6401,13 @@ impl GpuExperienceCollector {
|
||||
use crate::cuda_pipeline::sp14_isv_slots::HOLD_REWARD_EMA_INDEX;
|
||||
HOLD_REWARD_EMA_INDEX as i32
|
||||
})
|
||||
// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf
|
||||
// output for the replay-buffer schema. Threaded as
|
||||
// a `[N × L]` f32 buffer (counterfactual extension
|
||||
// happens at the `aux_conf` assembly site at the
|
||||
// end of `collect_experiences_gpu`, mirroring the
|
||||
// `aux_sign_labels` `[N × L × 2]` final shape).
|
||||
.arg(&mut self.aux_conf_per_sample)
|
||||
.launch(launch_cfg)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"experience_env_step t={t}: {e}"
|
||||
|
||||
@@ -73,7 +73,8 @@ fn insert_random_batch(
|
||||
let rewards = random_cuda_f32(n, stream)?;
|
||||
let dones = zeros_cuda_f32(n, stream)?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(n).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, n)?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(n).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, n)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -140,7 +140,8 @@ fn test_per_sample_latency() -> anyhow::Result<()> {
|
||||
let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, batch_insert)?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, batch_insert)?;
|
||||
}
|
||||
assert_eq!(buf.len(), fill_count);
|
||||
|
||||
|
||||
@@ -150,7 +150,8 @@ fn test_per_weights_valid() -> anyhow::Result<()> {
|
||||
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, 1)?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, 1)?;
|
||||
}
|
||||
|
||||
buf.sample_proportional(16)?;
|
||||
@@ -195,7 +196,8 @@ fn test_per_indices_valid() -> anyhow::Result<()> {
|
||||
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, 1)?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?;
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, 1)?;
|
||||
}
|
||||
|
||||
buf.sample_proportional(16)?;
|
||||
|
||||
@@ -2526,6 +2526,7 @@ impl DQNTrainer {
|
||||
&gpu_batch.rewards,
|
||||
&gpu_batch.dones,
|
||||
&gpu_batch.aux_sign_labels,
|
||||
&gpu_batch.aux_conf,
|
||||
total,
|
||||
).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?;
|
||||
}
|
||||
|
||||
@@ -122,8 +122,9 @@ fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) {
|
||||
let rf = stream.clone_htod(&rewards_host).unwrap();
|
||||
let df = stream.clone_htod(&dones_host).unwrap();
|
||||
let aux_sign: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
let aux_conf: CudaSlice<f32> = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, n).unwrap();
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, &aux_conf, n).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2165,8 +2165,9 @@ mod gpu {
|
||||
let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards");
|
||||
let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones");
|
||||
let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign");
|
||||
let ac: CudaSlice<f32> = stream.alloc_zeros::<f32>(n1).expect("aux_conf alloc");
|
||||
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, n1)
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, &ac, n1)
|
||||
.expect("insert_batch (recovery=1)");
|
||||
stream.synchronize().expect("sync after insert_batch (recovery=1)");
|
||||
|
||||
@@ -2183,8 +2184,9 @@ mod gpu {
|
||||
let rf2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod rf2");
|
||||
let df2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod df2");
|
||||
let ax2: CudaSlice<i32> = stream.clone_htod(&vec![0_i32; n2]).expect("htod ax2");
|
||||
let ac2: CudaSlice<f32> = stream.alloc_zeros::<f32>(n2).expect("aux_conf alloc 2");
|
||||
|
||||
buf.insert_batch(&sf2, &nsf2, &af2, &rf2, &df2, &ax2, n2)
|
||||
buf.insert_batch(&sf2, &nsf2, &af2, &rf2, &df2, &ax2, &ac2, n2)
|
||||
.expect("insert_batch (recovery=0)");
|
||||
stream.synchronize().expect("sync after insert_batch (recovery=0)");
|
||||
|
||||
@@ -2349,8 +2351,9 @@ mod gpu {
|
||||
let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards");
|
||||
let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones");
|
||||
let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign");
|
||||
let ac: CudaSlice<f32> = stream.alloc_zeros::<f32>(batch_size).expect("aux_conf alloc");
|
||||
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, batch_size)
|
||||
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, &ac, batch_size)
|
||||
.expect("insert_batch (mixed-env)");
|
||||
stream.synchronize().expect("sync after insert_batch");
|
||||
|
||||
|
||||
@@ -2,6 +2,194 @@
|
||||
|
||||
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
|
||||
|
||||
## 2026-05-10 — SP20 Phase 3 Task 3.4: replay-buffer schema for per-bar aux_conf (atomic)
|
||||
|
||||
Lands the producer→ring→consumer schema plumbing for per-bar
|
||||
`aux_conf` (the K=2 peak-softmax-above-uniform confidence value the
|
||||
Phase 5 Aux→Q gate consumes at the Bellman target). Atomic across all
|
||||
struct boundaries per `feedback_no_partial_refactor`:
|
||||
|
||||
```
|
||||
ExperienceCollector → GpuExperienceBatch → insert_batch() → ring → sample() → GpuBatchPtrs → Trainer
|
||||
(writes per (i,t)) (.aux_conf field) (8th arg) (ring (.aux_conf_ptr (Phase 5 gate
|
||||
column) field) consumer)
|
||||
```
|
||||
|
||||
Phase 5 lands the consumer (Bellman target gate kernel); this commit
|
||||
is the producer-side schema plumbing only.
|
||||
|
||||
### Spec §4.4 consumer contract (Phase 5 forward-reference)
|
||||
|
||||
```
|
||||
At each replay batch step:
|
||||
aux_conf = ptrs.aux_conf_ptr[k] // SAMPLED bar
|
||||
threshold = ISV[AUX_CONF_THRESHOLD_INDEX] // [0.01, 0.20]
|
||||
temp = ISV[AUX_GATE_TEMP_INDEX] // adaptive
|
||||
|
||||
gate = sigmoid((aux_conf - threshold) / temp) // ∈ [0, 1]
|
||||
|
||||
Q_target = gate × Q_target_full + (1 - gate) × Q_target_baseline
|
||||
where Q_target_baseline = mean_a Q(s, a)
|
||||
```
|
||||
|
||||
The gate ∈ [0, 1] gracefully degrades the Bellman target toward
|
||||
`mean_a Q(s, a)` at uncertain states (low aux_conf), preventing the
|
||||
all-actions punishment that drives Hold-everywhere. Spec §4.4 calls
|
||||
this the "uncertain-state neutralizer" semantic.
|
||||
|
||||
### Producer site (`experience_env_step` per (i, t))
|
||||
|
||||
New kernel arg appended at the end of the signature (after the Task
|
||||
3.2 `aux_logits_per_env` / `hold_baseline_buffer` / `hold_buffer_size`
|
||||
/ `aux_k` / `hold_cost_scale_idx` / `hold_reward_ema_idx` block):
|
||||
|
||||
```cuda
|
||||
float* __restrict__ aux_conf_per_sample /* [N × L] f32 */
|
||||
```
|
||||
|
||||
Written ONCE per (i, t) at the kernel-entry default block (just after
|
||||
the H6 diagnostic defaults at line ~2270), using the same
|
||||
`sp20_compute_per_bar_aux_conf_k2` helper Task 3.2 added in
|
||||
`sp20_hold_baseline.cuh`:
|
||||
|
||||
```cuda
|
||||
if (aux_conf_per_sample != NULL) {
|
||||
if (aux_logits_per_env != NULL) {
|
||||
const float* logits_row_init = aux_logits_per_env + i * aux_k;
|
||||
aux_conf_per_sample[out_off] = sp20_compute_per_bar_aux_conf_k2(logits_row_init);
|
||||
} else {
|
||||
aux_conf_per_sample[out_off] = 0.0f;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
NULL-tolerant: when `aux_logits_per_env == NULL` (test scaffolds), the
|
||||
slot retains the K=2 uniform-prior sentinel `0.0f` — no information.
|
||||
The Phase 5 gate at default 0.0 with threshold ∈ [0.01, 0.20] gives
|
||||
`gate ≈ 0` ⇒ Q_target falls back to `mean_a Q(s, a)` (graceful
|
||||
degradation).
|
||||
|
||||
### Counterfactual slot
|
||||
|
||||
The kernel ALSO writes `aux_conf_per_sample[cf_off]` inside the
|
||||
existing CF block (line ~4485-ish, after `out_actions[cf_off]` etc.):
|
||||
|
||||
```cuda
|
||||
if (aux_conf_per_sample != NULL) {
|
||||
aux_conf_per_sample[cf_off] = aux_conf_per_sample[out_off];
|
||||
}
|
||||
```
|
||||
|
||||
The CF state at `cf_off` shares the SAME env state (only the action
|
||||
differs), so it shares the SAME aux signal — mirror the on-policy
|
||||
write rather than recomputing.
|
||||
|
||||
### New collector buffer
|
||||
|
||||
`GpuExperienceCollector.aux_conf_per_sample: CudaSlice<f32>` —
|
||||
`[alloc_episodes × alloc_timesteps × cf_mult]` f32 row-major. Sized
|
||||
to match the kernel's `out_off` + `cf_off` layout (out_off ∈
|
||||
`[0, total_output)`, cf_off ∈ `[total_output, total_output × cf_mult)`).
|
||||
alloc_zeros default = K=2 uniform-prior sentinel.
|
||||
|
||||
### New batch struct field
|
||||
|
||||
`GpuExperienceBatch.aux_conf: CudaSlice<f32>` — `[total = base_total
|
||||
* cf_mult]` f32. Cloned from `aux_conf_per_sample` via
|
||||
`dtod_clone_f32_native(buf, total)` at the end of
|
||||
`collect_experiences_gpu` (mirrors states/rewards/actions/dones
|
||||
clone pattern).
|
||||
|
||||
### New `GpuBatchPtrs` field (`crates/ml-dqn/src/gpu_replay_buffer.rs`)
|
||||
|
||||
`GpuBatchPtrs.aux_conf_ptr: u64` — points at either the trainer's
|
||||
destination buffer (when `set_trainer_aux_conf_ptr` was called by
|
||||
Phase 5) or the fallback `sample_aux_conf` ring-gather destination
|
||||
(default).
|
||||
|
||||
### New `GpuReplayBuffer` ring + sample fields
|
||||
|
||||
- `aux_conf: CudaSlice<f32>` — `[capacity]` ring storage. Written
|
||||
by `scatter_insert_f32` in `insert_batch` (mirrors the
|
||||
`aux_sign_labels` i32 path immediately above; reuses the existing
|
||||
f32 scatter kernel that scatters rewards / dones).
|
||||
- `sample_aux_conf: CudaSlice<f32>` — `[max_batch_size]` per-batch
|
||||
gather destination. Written by `gather_f32` (fallback) or
|
||||
`gather_f32_scalar` (direct-to-trainer path).
|
||||
- `trainer_aux_conf_ptr: u64` — destination for the direct path
|
||||
(Phase 5 wires the trainer's `aux_conf_at_state_buf` via
|
||||
`set_trainer_aux_conf_ptr`).
|
||||
|
||||
### `insert_batch` signature change (atomic across all callers)
|
||||
|
||||
Added 7th arg `aux_conf_in: &CudaSlice<f32>` between `aux_sign` and
|
||||
`bs`. All 6 callers updated atomically:
|
||||
|
||||
- `crates/ml/src/trainers/dqn/trainer/training_loop.rs:2522` (production)
|
||||
- `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs:77`
|
||||
- `crates/ml/src/trainers/dqn/smoke_tests/performance.rs:144`
|
||||
- `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs:154 + :201`
|
||||
- `crates/ml/tests/gpu_per_integration_test.rs:128`
|
||||
- `crates/ml/tests/sp15_phase1_oracle_tests.rs:2170 + :2189 + :2356`
|
||||
- 3 inline tests in `gpu_replay_buffer.rs::tests` (1528, 1608, 1672)
|
||||
|
||||
Each smoke / oracle test scaffold passes a fresh `alloc_zeros::<f32>(n)`
|
||||
for aux_conf — matches the K=2 uniform-prior semantic when the aux
|
||||
head isn't exercised by that scaffold.
|
||||
|
||||
### Tests
|
||||
|
||||
- `gpu_replay_buffer::tests::test_aux_conf_schema_round_trip` (GPU
|
||||
behavioral test): inserts 8 transitions with distinct in-range
|
||||
aux_conf values [0.0, 0.05, ..., 0.35]; samples 1; asserts the
|
||||
gathered slot value matches at least one inserted value
|
||||
(round-trip integrity invariant per
|
||||
`pearl_tests_must_prove_not_lock_observations`).
|
||||
- `gpu_replay_buffer::tests::test_aux_conf_trainer_ptr_wiring_contract`
|
||||
(CPU-only): asserts `set_trainer_aux_conf_ptr` /
|
||||
`trainer_aux_conf_ptr` setter/accessor pair behavior matches
|
||||
`set_isv_signals_ptr` / `isv_signals_dev_ptr` mirror.
|
||||
|
||||
### Files modified
|
||||
|
||||
| File | Status | Purpose |
|
||||
|------|--------|---------|
|
||||
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | +1 kernel arg, +2 write sites | Per-bar aux_conf production |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | +field, +alloc, +arg pass, +clone, +batch field | Producer-side schema |
|
||||
| `crates/ml-dqn/src/gpu_replay_buffer.rs` | +`GpuBatchPtrs.aux_conf_ptr`, +ring/sample fields, +scatter/gather, +setter/accessor, +insert_batch arg | Replay buffer schema |
|
||||
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | +1 arg in insert_batch call | Production caller |
|
||||
| `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs` | +1 arg | Smoke test |
|
||||
| `crates/ml/src/trainers/dqn/smoke_tests/performance.rs` | +1 arg | Smoke test |
|
||||
| `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | +1 arg ×2 | Smoke tests |
|
||||
| `crates/ml/tests/gpu_per_integration_test.rs` | +1 arg | Integration test |
|
||||
| `crates/ml/tests/sp15_phase1_oracle_tests.rs` | +1 arg ×3 | Oracle tests |
|
||||
| `docs/dqn-wire-up-audit.md` | This entry | Audit log |
|
||||
|
||||
### Verification
|
||||
|
||||
```
|
||||
SQLX_OFFLINE=true cargo check -p ml --tests --features cuda # green
|
||||
SQLX_OFFLINE=true cargo check -p ml-dqn --tests --features cuda # green
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib sp20 # 21/21 pass
|
||||
SQLX_OFFLINE=true cargo test -p ml-dqn test_aux_conf_trainer_ptr # CPU pass
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml-dqn --features cuda \
|
||||
test_aux_conf_schema_round_trip # GPU host pass
|
||||
```
|
||||
|
||||
### Plan accuracy errata
|
||||
|
||||
See `docs/superpowers/plans/2026-05-09-sp19-20-wr-first.md` Phase 3
|
||||
Errata Gap 11 for the `GpuBatchPtrs` vs `GpuExperienceBatch` naming
|
||||
clarification (both structs got the new field; the data flow has TWO
|
||||
struct boundaries, not one).
|
||||
|
||||
### Phase 3 → Phase 5 forward references
|
||||
|
||||
- Phase 5: Aux→Q gate kernel reads `GpuBatchPtrs.aux_conf_ptr` at
|
||||
the Bellman target loss site. `set_trainer_aux_conf_ptr` is
|
||||
called once after the trainer's `aux_conf_at_state_buf` is
|
||||
allocated, switching the gather to direct-to-trainer mode.
|
||||
|
||||
## 2026-05-10 — SP20 Phase 3 Task 3.2: Hold opportunity-cost dual emission (Component 2, atomic)
|
||||
|
||||
Lands the per-bar Hold opportunity-cost dual emission at the
|
||||
|
||||
Reference in New Issue
Block a user