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

feat(rl): Phase 4.3 — V_dq → PPO advantage swap + target net soft-update

Browse Source 
Activates DuelingQHead's V output as the PPO advantage baseline,
replacing scalar value_head's contribution. Also wires soft-update
of DuelingQHead's target net (was deferred from 4.2).

Two changes:

1. DuelingQHead.soft_update_target(): reuses dqn_target_soft_update
   kernel (generic element-wise blend with τ from ISV[401]) across
   all 4 weight tensors (w_v, b_v, w_a, b_a). Called once per training
   step after Adam, mirroring DQN's pattern.

2. compute_advantage_return call: v_pred_d → dueling_v_d,
   v_pred_tp1_d → dueling_v_tp1_d. PPO advantage is now:
       A(s_t, a_t) = E_Q_C51(s_t, a_taken) − V_dq(s_t)
       returns(s_t) = r_t + γ(1−done) × V_dq(s_{t+1})
   value_head still trains via MSE on returns for diagnostic
   comparison; can be retired in a future commit if V_dq proves
   itself.

Phase 4.2 validated structurally at b=1024 5k that DuelingQHead's
training does not perturb Plan A v2's dynamics (qpa tracked within
0.02 at every milestone). Phase 4.3 is the smallest possible commit
that USES V_dq downstream — fully de-risked by 4.2's structural test.

Cluster expectation: qpa trajectory roughly matches Plan A v2 (V_dq
calibrated by joint training on same Bellman reward signal as C51,
so cross-architecture scale issue from Phase 3.2 doesn't apply here).

Validated:
  - cargo build --release clean
  - integrated_trainer_smoke 1 step passes
  - alpha_rl_train --steps 3 --b 128 under compute-sanitizer: 0 errors,
    l_q=0.024, l_v=0.0003 (slightly higher than Plan A v2's 0.0001
    because V_scalar now sees different gradient prop now that V_dq
    drives advantage — still healthy)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-30 09:14:46 +02:00
parent acdafe508e
commit 25f5ce99b6
2 changed files with 68 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
crates/ml-alpha/src/rl/dueling_q.rs
View File

@@ -73,6 +73,13 @@ const DUELING_Q_DECOMPOSE_AND_BWD_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/rl_dueling_q_decompose_and_bwd.cubin"
));
/// Reuse the existing dqn_target_soft_update kernel — it's a generic
/// element-wise `target = (1−τ) target + τ online` blend that works
/// for any tensor size. τ read from ISV[RL_TARGET_TAU_INDEX=401].
const DQN_TARGET_SOFT_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/dqn_target_soft_update.cubin"
));
/// Construction config for [`DuelingQHead`].
#[derive(Clone, Debug)]
@@ -112,6 +119,8 @@ pub struct DuelingQHead {
pub loss_and_grad_fn: CudaFunction,
_bwd_module: Arc<CudaModule>,
pub decompose_and_bwd_fn: CudaFunction,
_soft_update_module: Arc<CudaModule>,
pub soft_update_fn: CudaFunction,
// ── Online weights ───────────────────────────────────────────────
/// V projection weights `[HIDDEN_DIM]`.
@@ -165,6 +174,13 @@ impl DuelingQHead {
let decompose_and_bwd_fn = bwd_module
.load_function("rl_dueling_q_decompose_and_weight_grad")
.context("load rl_dueling_q_decompose_and_weight_grad")?;
// Reuse dqn_target_soft_update kernel — generic element-wise blend.
let soft_update_module = ctx
.load_cubin(DQN_TARGET_SOFT_UPDATE_CUBIN.to_vec())
.context("load dqn_target_soft_update cubin for dueling")?;
let soft_update_fn = soft_update_module
.load_function("dqn_target_soft_update")
.context("load dqn_target_soft_update for dueling")?;
// ── Weight init (Xavier uniform, scaled by 0.01 like sibling heads) ──
// Per pearl_scoped_init_seed_for_reproducibility.
@@ -212,6 +228,8 @@ impl DuelingQHead {
loss_and_grad_fn,
_bwd_module: bwd_module,
decompose_and_bwd_fn,
_soft_update_module: soft_update_module,
soft_update_fn,
w_v_d,
b_v_d,
w_a_d,
@@ -450,6 +468,39 @@ impl DuelingQHead {
}
Ok(())
}
/// Soft-update target network: `target = (1−τ) × target + τ × online`
/// element-wise on all four weight tensors. τ read from
/// `ISV[RL_TARGET_TAU_INDEX=401]` — same controller as Q's target.
/// Should be called once per training step AFTER the Adam steps so
/// the update reflects the latest online weights.
pub fn soft_update_target(&mut self, isv_dev_ptr: &u64) -> Result<()> {
let launch_blend = |n: usize, online: u64, target: u64, label: &str| -> Result<()> {
let n_i = n as i32;
let mut args = RawArgs::new();
args.push_ptr(online);
args.push_ptr(target);
args.push_ptr(*isv_dev_ptr);
args.push_i32(n_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.soft_update_fn.cu_function(),
(((n as u32) + 255) / 256, 1, 1),
(256, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dueling soft_update_target ({label}): {e:?}"))?;
}
Ok(())
};
launch_blend(self.w_v_d.len(), self.w_v_d.raw_ptr(), self.w_v_target_d.raw_ptr(), "w_v")?;
launch_blend(self.b_v_d.len(), self.b_v_d.raw_ptr(), self.b_v_target_d.raw_ptr(), "b_v")?;
launch_blend(self.w_a_d.len(), self.w_a_d.raw_ptr(), self.w_a_target_d.raw_ptr(), "w_a")?;
launch_blend(self.b_a_d.len(), self.b_a_d.raw_ptr(), self.b_a_target_d.raw_ptr(), "b_a")?;
Ok(())
}
}
// ── pinned-staging upload helper (mirrors aux_heads.rs::upload) ──

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

@@ -6418,14 +6418,23 @@ impl IntegratedTrainer {
// calibration directly into the PPO advantage signal, fixing
// the q_pi_agree anti-correlation where PPO with V-advantage
// pushed π away from Q's preferred actions.
//
// Phase 4.3 (2026-05-30): V baseline now sourced from
// DuelingQHead's dueling-trained V output (dueling_v_d /
// dueling_v_tp1_d) — replaces scalar value_head's v_pred_d.
// The dueling architecture's V is trained jointly with A via
// Bellman loss on composed Q, providing an explicitly
// V/A-decomposed baseline. value_head still trains via MSE on
// returns for diagnostic comparison (V_dq vs V_scalar tracking).
// Per spec docs/superpowers/specs/2026-05-30-phase4-independent-dueling-head-design.md §6.
{
let grid_x = ((b_size as u32) + 31) / 32;
let mut args = RawArgs::new();
args.push_ptr(self.isv_dev_ptr);
args.push_ptr(self.rewards_d.raw_ptr());
args.push_ptr(self.dones_d.raw_ptr());
args.push_ptr(self.v_pred_d.raw_ptr());
args.push_ptr(self.v_pred_tp1_d.raw_ptr());
args.push_ptr(self.dueling_v_d.raw_ptr());
args.push_ptr(self.dueling_v_tp1_d.raw_ptr());
args.push_ptr(self.returns_d.raw_ptr());
args.push_ptr(self.advantages_d.raw_ptr());
args.push_i32(b_size_i);
@@ -7363,6 +7372,12 @@ impl IntegratedTrainer {
.step(&mut self.dueling_q_head.b_a_d, &self.dueling_grad_b_a_d)
.context("dueling_q_b_a_adam.step")?;
// Phase 4.3: soft-update target net AFTER Adam (target absorbs
// latest online weights with τ from ISV[RL_TARGET_TAU_INDEX]).
self.dueling_q_head
.soft_update_target(&self.isv_dev_ptr)
.context("dueling_q_head.soft_update_target")?;
// Ensemble Q
{
let b_size_i = b_size as i32;