diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 7473190a3..b7c8a5fd5 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -91,6 +91,7 @@ use rand::{Rng, SeedableRng}; use crate::rl::common::{N_ACTIONS, Q_N_ATOMS, Q_V_MAX, Q_V_MIN}; use crate::rl::dqn::{DqnHead, DqnHeadConfig}; use crate::rl::iqn::EMBED_DIM; +use crate::rl::noisy::{NoisyLinear, NoisyLinearConfig}; use crate::rl::isv_slots::{ RL_LR_BCE_INDEX, RL_LR_AUX_INDEX, RL_LR_PI_INDEX, RL_LR_Q_INDEX, RL_LR_V_INDEX, RL_SLOTS_END, @@ -225,6 +226,13 @@ const RL_TRAIL_STOP_CHECK_CUBIN: &[u8] = const ACTIONS_TO_MARKET_TARGETS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/actions_to_market_targets.cubin")); +// Element-wise in-place add: dst[i] += src[i]. Used by noisy +// exploration to fold NoisyLinear output into ensemble_q_d before +// action selection. Loaded from the same cubin that perception.rs +// uses for aux→encoder gradient accumulation. +const AUX_VEC_ADD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_vec_add.cubin")); + // Phase R7a: element-wise abs-copy kernel for the |reward| signal // feeding ema_update_on_done's MEAN_ABS_PNL_EMA slot. const ABS_COPY_CUBIN: &[u8] = @@ -669,6 +677,32 @@ pub struct IntegratedTrainer { /// Host RNG for tau sampling (deterministic from dqn_seed + offset). iqn_tau_rng: rand_chacha::ChaCha8Rng, + // ── NoisyNet state-dependent exploration ───────────────────────── + // Factored noisy linear layer (Fortunato et al. 2017) that maps + // h_t → [B, N_ACTIONS] state-dependent exploration noise. Added + // to ensemble_q_d before action selection so the policy sees + // noisy Q values during rollout. Replaces the P_MIN probability + // floor in rl_pi_action_kernel. Target network stays deterministic + // (noise buffers zeroed = no resample_noise calls on target path). + + /// NoisyLinear layer: h_t [B, HIDDEN_DIM] → noise [B, N_ACTIONS]. + /// Resampled once per step_with_lobsim via resample_noise(). + pub noisy_exploration: NoisyLinear, + /// Output buffer for the NoisyLinear forward [B × N_ACTIONS]. + pub noisy_output_d: CudaSlice, + /// Adam optimiser for NoisyLinear mu_w [N_ACTIONS × HIDDEN_DIM]. + pub noisy_mu_w_adam: AdamW, + /// Adam optimiser for NoisyLinear sigma_w [N_ACTIONS × HIDDEN_DIM]. + pub noisy_sigma_w_adam: AdamW, + /// Adam optimiser for NoisyLinear mu_b [N_ACTIONS]. + pub noisy_mu_b_adam: AdamW, + /// Adam optimiser for NoisyLinear sigma_b [N_ACTIONS]. + pub noisy_sigma_b_adam: AdamW, + /// Kernel handle for `aux_vec_add_inplace` — folds noisy_output_d + /// into ensemble_q_d before the π action sampler reads it. + _aux_vec_add_module: Arc, + aux_vec_add_fn: CudaFunction, + // ── Phase R7d: PER + off-policy Q replay ────────────────────────── /// Prioritised Experience Replay buffer (off-policy DQN training). /// Per `feedback_always_per` PER is always-on. Wired in R7d after @@ -1395,6 +1429,47 @@ impl IntegratedTrainer { cfg.dqn_seed.wrapping_add(0x1C_A0_7A_00), ); + // NoisyNet exploration layer: h_t → [B, N_ACTIONS] noise added + // to ensemble_q_d before action selection. Replaces the P_MIN + // probability floor with learned, state-dependent exploration. + // Per pearl_scoped_init_seed_for_reproducibility: seed offset + // +200 from dqn_seed keeps Xavier draws independent of C51/IQN. + let noisy_exploration = NoisyLinear::new( + dev, + NoisyLinearConfig { + in_dim: HIDDEN_DIM, + out_dim: N_ACTIONS, + sigma_init: 0.5, // canonical NoisyNet; ISV slot 547 overrides at step time + seed: cfg.dqn_seed.wrapping_add(200), + }, + ) + .context("NoisyLinear::new (exploration)")?; + let noisy_output_d = stream + .alloc_zeros::(b_size * N_ACTIONS) + .context("alloc noisy_output_d")?; + // Adam state for the 4 NoisyLinear weight tensors. LR mirrored + // from ISV[RL_LR_Q_INDEX] (same as the Q head — the noisy layer + // is part of the Q exploration pipeline). + let noisy_mu_w_adam = + AdamW::new(dev, noisy_exploration.mu_w.len(), lr_placeholder) + .context("noisy_mu_w_adam")?; + let noisy_sigma_w_adam = + AdamW::new(dev, noisy_exploration.sigma_w.len(), lr_placeholder) + .context("noisy_sigma_w_adam")?; + let noisy_mu_b_adam = + AdamW::new(dev, noisy_exploration.mu_b.len(), lr_placeholder) + .context("noisy_mu_b_adam")?; + let noisy_sigma_b_adam = + AdamW::new(dev, noisy_exploration.sigma_b.len(), lr_placeholder) + .context("noisy_sigma_b_adam")?; + // aux_vec_add_inplace for folding noisy output into ensemble_q_d. + let aux_vec_add_module = ctx + .load_cubin(AUX_VEC_ADD_CUBIN.to_vec()) + .context("load aux_vec_add cubin")?; + let aux_vec_add_fn = aux_vec_add_module + .load_function("aux_vec_add_inplace") + .context("load aux_vec_add_inplace fn")?; + // Phase R7d: PER buffer + sampled-batch device scratch. // The replay buffer holds per-transition device payloads // (h_t / next_h_t as small `CudaSlice(HIDDEN_DIM)` slices @@ -1624,6 +1699,14 @@ impl IntegratedTrainer { iqn_expected_q_d, ensemble_q_d, iqn_tau_rng, + noisy_exploration, + noisy_output_d, + noisy_mu_w_adam, + noisy_sigma_w_adam, + noisy_mu_b_adam, + noisy_sigma_b_adam, + _aux_vec_add_module: aux_vec_add_module, + aux_vec_add_fn, replay, n_step_buffer: (0..b_size).map(|_| Vec::with_capacity(16)).collect(), sampled_h_t_d, @@ -1709,7 +1792,7 @@ impl IntegratedTrainer { // (slot, value) pair — pure device write, no HtoD per // `feedback_no_htod_htoh_only_mapped_pinned`. { - let isv_constants: [(usize, f32); 79] = [ + let isv_constants: [(usize, f32); 83] = [ // Static seeds for the adaptive reward-clamp controller — // these are the initial values that // `rl_reward_clamp_controller` will replace once it @@ -1775,7 +1858,13 @@ impl IntegratedTrainer { (crate::rl::isv_slots::RL_CONF_GATE_SIGMA_NORM_INDEX, 1.0), (crate::rl::isv_slots::RL_FRD_GATE_THR_LONG_INDEX, 0.05), (crate::rl::isv_slots::RL_FRD_GATE_THR_SHORT_INDEX, 0.05), - (crate::rl::isv_slots::RL_PI_SAMPLE_P_MIN_INDEX, 0.015), + // P_MIN set to 0.0: NoisyNet state-dependent exploration + // replaces the fixed probability floor. Per + // `pearl_pi_actor_collapses_without_entropy_floor`: + // NoisyLinear noise on ensemble_q_d provides learned + // exploration, so the P_MIN anti-collapse floor in + // rl_pi_action_kernel is no longer needed. + (crate::rl::isv_slots::RL_PI_SAMPLE_P_MIN_INDEX, 0.0), (crate::rl::isv_slots::RL_OUTCOME_ALPHA_INDEX, 0.1), (crate::rl::isv_slots::RL_GATE_WARMUP_STEPS_INDEX, 10000.0), (crate::rl::isv_slots::RL_GATE_DONES_TARGET_INDEX, 0.10), @@ -1823,6 +1912,16 @@ impl IntegratedTrainer { (crate::rl::isv_slots::RL_ROLLOUT_BOOTSTRAP_INDEX, 2048.0), (crate::rl::isv_slots::RL_REWARD_SCALE_BOOTSTRAP_INDEX, 1.0), (crate::rl::isv_slots::RL_PPO_RATIO_CLAMP_BOOTSTRAP_INDEX, 10.0), + // IQN ensemble seeds — alpha=0.5 (equal C51/IQN blend), + // N_TAU=32, LR=1e-3 (canonical). + (crate::rl::isv_slots::RL_IQN_N_TAU_INDEX, 32.0), + (crate::rl::isv_slots::RL_IQN_ENSEMBLE_ALPHA_INDEX, 0.5), + (crate::rl::isv_slots::RL_IQN_LR_INDEX, 1e-3), + // NoisyNet σ_init — canonical Fortunato et al. 2017 value. + // Read by NoisyLinear at construction (fallback if ISV is + // sentinel-zero) and available for runtime re-tuning via + // diagnostic JSONL visibility. + (crate::rl::isv_slots::RL_NOISY_SIGMA_INIT_INDEX, 0.5), ]; let cfg_isv = LaunchConfig { grid_dim: (1, 1, 1), @@ -4026,6 +4125,43 @@ impl IntegratedTrainer { } } + // ── NoisyNet state-dependent exploration on ensemble Q ───────── + // Resample factored noise, forward h_t through the NoisyLinear + // layer → noisy_output_d [B, N_ACTIONS], then fold into + // ensemble_q_d in-place via aux_vec_add_inplace. The perturbed + // ensemble Q feeds the Q→π agreement diagnostic and distillation + // gradient, so the policy indirectly sees noisy Q preferences + // and learns to explore without a fixed P_MIN floor. + // + // Target network path (Bellman argmax / bellman_target_projection) + // does NOT see noise — it reads the un-perturbed C51/IQN target + // weights. This matches the NoisyNet paper's prescription: + // noise on the online network only during rollout. + self.noisy_exploration + .resample_noise() + .context("step_with_lobsim: noisy_exploration.resample_noise")?; + self.noisy_exploration + .forward(h_t_borrow, b_size, &mut self.noisy_output_d) + .context("step_with_lobsim: noisy_exploration.forward")?; + { + let n_elems = (b_size * N_ACTIONS) as i32; + let cfg_add = LaunchConfig { + grid_dim: (((n_elems as u32) + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.aux_vec_add_fn); + launch + .arg(&mut self.ensemble_q_d) + .arg(&self.noisy_output_d) + .arg(&n_elems); + unsafe { + launch + .launch(cfg_add) + .context("step_with_lobsim: aux_vec_add_inplace (noisy → ensemble_q)")?; + } + } + // ── Step 2b: Forward π logits for log_pi_old. ───────────────── let mut pi_logits_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; self.policy_head