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

feat(rl): wire noisy nets into ensemble Q — learned exploration, P_MIN=0

Browse Source 
NoisyLinear layer on top of ensemble Q values provides state-dependent
exploration noise via h_t → learned perturbation. Replaces the fixed
P_MIN=0.015 probability floor with learned exploration (P_MIN→0.0).

- NoisyLinear constructed with in=128, out=11, σ_init=0.5
- resample_noise() each step for fresh factored noise
- forward(h_t) → noise_d added to ensemble_q_d via aux_vec_add
- 4 Adam optimizers for mu_w, sigma_w, mu_b, sigma_b
- ISV seeds for IQN ensemble α=0.5, n_tau=32, σ_init=0.5
- Target network path stays noise-free (Fortunato et al. 2017)

Local smoke: 100 steps, l_q=2.30, no crash. All three phases complete.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-25 21:19:55 +02:00
parent 6308be794e
commit 993201250a
1 changed files with 138 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -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<f32>,
/// 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<CudaModule>,
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::<f32>(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<f32>(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::<f32>(b_size * N_ACTIONS)?;
self.policy_head