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) <noreply@anthropic.com>
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3976,6 +3976,9 @@ impl GpuDqnTrainer {
|
||||
pub fn td_errors_buf(&self) -> &CudaSlice<f32> {
|
||||
&self.td_errors_buf
|
||||
}
|
||||
pub fn td_errors_buf_mut(&mut self) -> &mut CudaSlice<f32> {
|
||||
&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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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() }
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user