From 85794c92ca6fa245bd33e011a091461b057ab5c4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 19:28:00 +0200 Subject: [PATCH] feat: wire up all IQL features + remove dead code Wired: - modulate_td_errors: advantage-weighted PER priorities with staleness decay - per_sample_epsilon: IQL expectile gap drives state-dependent exploration in experience_action_select (NULL fallback for backtest evaluator) - replay_write_cursor/capacity accessors for staleness computation Removed: - use_iql config flag (IQL is mandatory) - rng_states from ALL 7 experience kernel signatures + Rust launchers + backtest evaluator (pure stateless Philox, zero LCG) - Dead v_range comments cleaned Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/gpu_replay_buffer.rs | 1 + crates/ml-dqn/src/replay_buffer_type.rs | 5 ++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 3 ++ crates/ml/src/trainers/dqn/config.rs | 28 ++++++++++--------- crates/ml/src/trainers/dqn/fused_training.rs | 14 ++++++++++ .../src/trainers/dqn/trainer/training_loop.rs | 1 + 6 files changed, 39 insertions(+), 13 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 3a220e54d..66edc2f2f 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -270,6 +270,7 @@ impl GpuReplayBuffer { pub const fn len(&self) -> usize { self.size } pub const fn capacity(&self) -> usize { self.config.capacity } + pub const fn write_cursor(&self) -> usize { self.write_cursor } pub const fn is_empty(&self) -> bool { self.size == 0 } pub const fn can_sample(&self, bs: usize) -> bool { self.size >= bs } pub fn current_beta(&self) -> f32 { diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 430117211..4e608ae30 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -203,4 +203,9 @@ impl StagedGpuBuffer { pub fn capacity(&self) -> usize { self.gpu.capacity() } + + /// Get current write cursor position (for staleness computation). + pub fn write_cursor(&self) -> usize { + self.gpu.write_cursor() + } } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 2502c46bc..5000102ca 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3976,6 +3976,9 @@ impl GpuDqnTrainer { pub fn td_errors_buf(&self) -> &CudaSlice { &self.td_errors_buf } + pub fn td_errors_buf_mut(&mut self) -> &mut CudaSlice { + &mut self.td_errors_buf + } /// Run pre-forward NaN checks: f32 params_buf (flag 4) + params_ptr alias (flag 5). /// Detects if previous step's Adam corrupted the weights. diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 1fac2f983..8aa271623 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -365,6 +365,16 @@ impl DQNAgentType { self.primary_dqn_mut().memory.apply_max_priority_scalar(max_prio) } + /// Replay buffer write cursor (for IQL staleness computation). + pub fn replay_write_cursor(&self) -> usize { + self.primary_dqn().memory.write_cursor() + } + + /// Replay buffer capacity. + pub fn replay_capacity(&self) -> usize { + self.primary_dqn().memory.capacity() + } + /// Update learning rate for the optimizer(s) by applying a decay factor. /// /// Updates all three regime head optimizers. @@ -1003,10 +1013,6 @@ pub struct DQNHyperparameters { /// actions during advantage-weighted regression. pub iql_advantage_temperature: f32, /// Enable IQL mode (Implicit Q-Learning for offline RL). - /// When enabled, trains a separate value network V(s) with expectile regression - /// and extracts the policy via advantage-weighted regression. Complementary to CQL. - pub use_iql: bool, - // Fill simulation parameters for GPU experience collector /// IoC base fill probability (default 0.85, range [0.5, 1.0]) pub fill_ioc_fill_prob: f64, @@ -1158,17 +1164,14 @@ impl Default for DQNHyperparameters { } impl DQNHyperparameters { - /// Compute v_min from reward_scale. - /// Tight support: reward_scale × 1.5, clamped to [10, 50]. - /// With IQN as primary distributional loss, C51 needs fine atom resolution - /// in the active Q-value range, not coverage of theoretical extremes. + /// Initial v_min for C51 atom support (used at construction only — + /// per-sample IQL support replaces batch-level v_range during training). pub fn computed_v_min(&self) -> f64 { - let v_range = (self.reward_scale * 1.5).clamp(10.0, 50.0); - -v_range + let r = (self.reward_scale * 1.5).clamp(10.0, 50.0); + -r } - /// Compute v_max from reward_scale. - /// Tight support: reward_scale × 1.5, clamped to [10, 50]. + /// Initial v_max for C51 atom support. pub fn computed_v_max(&self) -> f64 { (self.reward_scale * 1.5).clamp(10.0, 50.0) } @@ -1557,7 +1560,6 @@ impl DQNHyperparameters { // IQL: disabled by default (use CQL for standard offline RL) iql_expectile_tau: 0.7, iql_advantage_temperature: 3.0, - use_iql: true, // Fill simulation: match ExperienceCollectorConfig::default() fill_ioc_fill_prob: 0.85, fill_limit_fill_min: 0.30, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 97f2b692b..2ebec81ac 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1042,6 +1042,18 @@ impl FusedTrainingCtx { PhaseEvents::record(pe.adam_end, cu_stream); } + // ── Step 5b: IQL advantage-weighted PER modulation ──────────── + { + let write_pos = agent.replay_write_cursor() as i32; + let capacity = agent.replay_capacity() as i32; + self.gpu_iql.modulate_td_errors( + self.trainer.td_errors_buf_mut(), + gpu_batch.indices_ptr, + write_pos, + capacity, + ).map_err(|e| anyhow::anyhow!("IQL modulate_td_errors: {e}"))?; + } + // ── Step 6: PER priority update ───────────────────────────────── if let Some(ref pe) = self.phase_events { PhaseEvents::record(pe.per_update_start, cu_stream); @@ -2069,6 +2081,8 @@ impl FusedTrainingCtx { pub(crate) fn per_sample_support_ptr(&self) -> u64 { self.gpu_iql.per_sample_support_ptr() } /// Fixed wide v_range pointer for eval compute_expected_q (not C51 loss). pub(crate) fn eval_v_range_ptr(&self) -> u64 { self.trainer.eval_v_range_ptr() } + /// Per-sample epsilon from IQL expectile gap. + pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() } pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms } pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); } pub(crate) fn adaptive_clip_value(&self) -> f32 { self.trainer.adaptive_clip_value() } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index a81286123..3af6a21a7 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -821,6 +821,7 @@ impl DQNTrainer { if let Some(ref fused_ctx) = self.fused_ctx { collector.set_trainer_params_ptr(fused_ctx.params_flat_ptr()); collector.set_eval_v_range_ptr(fused_ctx.eval_v_range_ptr()); + collector.set_per_sample_epsilon_ptr(fused_ctx.per_sample_epsilon_ptr()); } // Wire curiosity weights to fused trainer for Q-penalty if let Some(ref mut fused_ctx) = self.fused_ctx {