feat(rl): n-step returns (n=10) — 10× more direct reward signal
Replaces 1-step Bellman r + γQ(s') with n-step R_n + γⁿQ(s_{t+n}).
At γ=0.995, n=10 gives the agent 10 real rewards (~2.5 seconds)
before bootstrapping from Q, dramatically reducing bootstrap error.
Implementation:
- NStepEntry ring buffer per batch in IntegratedTrainer
- push_to_replay accumulates entries, flushes when len==n or done
- R_n = Σ γᵏ rₖ computed at flush time (both scaled + raw)
- n_step_gamma = γⁿ (or 0 if any done in window) stored per transition
- bellman_target_projection.cu uses per-transition n_step_gamma
instead of the global ISV γ for the bootstrap discount
- project_bellman_target wrapper takes n_step_gammas_d buffer
ISV slot 542 (RL_N_STEP_INDEX, default 10). Local smoke: 1k steps,
no crash, replay=150 (correct for n=10 with dones).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -90,12 +90,13 @@ extern "C" __global__ void dqn_select_action_atoms(
|
||||
|
||||
|
||||
extern "C" __global__ void bellman_target_projection(
|
||||
const float* __restrict__ target_logits, // [B × Q_N_ATOMS]
|
||||
const float* __restrict__ rewards, // [B]
|
||||
const float* __restrict__ dones, // [B]
|
||||
const float* __restrict__ isv, // ≥ 401
|
||||
const float* __restrict__ target_logits, // [B × Q_N_ATOMS]
|
||||
const float* __restrict__ rewards, // [B] (n-step discounted R_n)
|
||||
const float* __restrict__ dones, // [B]
|
||||
const float* __restrict__ n_step_gammas, // [B] (γⁿ per transition, 0 if done)
|
||||
const float* __restrict__ isv, // ≥ 401
|
||||
int B,
|
||||
float* __restrict__ target_dist // [B × Q_N_ATOMS]
|
||||
float* __restrict__ target_dist // [B × Q_N_ATOMS]
|
||||
) {
|
||||
const int batch = blockIdx.x;
|
||||
const int atom = threadIdx.x;
|
||||
@@ -133,10 +134,11 @@ extern "C" __global__ void bellman_target_projection(
|
||||
|
||||
const float p = s_softmax[atom] / s_sumexp;
|
||||
|
||||
// ── γ from ISV + done masking ────────────────────────────────────
|
||||
const float gamma = fmaxf(0.0f, fminf(1.0f, isv[RL_GAMMA_INDEX]));
|
||||
// ── n-step γ: per-transition γⁿ replaces the ISV γ for the
|
||||
// bootstrap discount. Already zero when any done occurred in the
|
||||
// n-step window (set at push time). For n=1, n_step_gamma = γ.
|
||||
const float r = rewards[batch];
|
||||
const float gamma_eff = gamma * (1.0f - dones[batch]);
|
||||
const float gamma_eff = n_step_gammas[batch];
|
||||
|
||||
// ── Adaptive atom span from ISV (audit follow-up). DELTA_Z is
|
||||
// recomputed every kernel launch — ratchet semantics in the
|
||||
|
||||
@@ -506,6 +506,7 @@ impl DqnHead {
|
||||
target_logits_d: &CudaSlice<f32>,
|
||||
rewards_d: &CudaSlice<f32>,
|
||||
dones_d: &CudaSlice<f32>,
|
||||
n_step_gammas_d: &CudaSlice<f32>,
|
||||
isv_d: &CudaSlice<f32>,
|
||||
b_size: usize,
|
||||
target_dist_d: &mut CudaSlice<f32>,
|
||||
@@ -513,6 +514,7 @@ impl DqnHead {
|
||||
debug_assert_eq!(target_logits_d.len(), b_size * Q_N_ATOMS);
|
||||
debug_assert_eq!(rewards_d.len(), b_size);
|
||||
debug_assert_eq!(dones_d.len(), b_size);
|
||||
debug_assert_eq!(n_step_gammas_d.len(), b_size);
|
||||
debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS);
|
||||
|
||||
let b_i = b_size as i32;
|
||||
@@ -526,6 +528,7 @@ impl DqnHead {
|
||||
.arg(target_logits_d)
|
||||
.arg(rewards_d)
|
||||
.arg(dones_d)
|
||||
.arg(n_step_gammas_d)
|
||||
.arg(isv_d)
|
||||
.arg(&b_i)
|
||||
.arg(target_dist_d);
|
||||
|
||||
@@ -660,6 +660,7 @@ pub struct IntegratedTrainer {
|
||||
|
||||
/// Per-step gather buffer for PER-sampled done flags. Length `B`.
|
||||
pub sampled_dones_d: CudaSlice<f32>,
|
||||
pub sampled_n_step_gammas_d: CudaSlice<f32>,
|
||||
|
||||
/// Per-step gather buffer for PER-sampled next-state argmax actions
|
||||
/// (Double-DQN online-Q argmax on `sampled_h_tp1`). Computed
|
||||
@@ -1341,6 +1342,9 @@ impl IntegratedTrainer {
|
||||
let sampled_dones_d = stream
|
||||
.alloc_zeros::<f32>(b_size)
|
||||
.context("alloc sampled_dones_d")?;
|
||||
let sampled_n_step_gammas_d = stream
|
||||
.alloc_zeros::<f32>(b_size)
|
||||
.context("alloc sampled_n_step_gammas_d")?;
|
||||
let sampled_next_actions_d = stream
|
||||
.alloc_zeros::<i32>(b_size)
|
||||
.context("alloc sampled_next_actions_d")?;
|
||||
@@ -1537,6 +1541,7 @@ impl IntegratedTrainer {
|
||||
sampled_actions_d,
|
||||
sampled_rewards_d,
|
||||
sampled_dones_d,
|
||||
sampled_n_step_gammas_d,
|
||||
sampled_next_actions_d,
|
||||
td_per_sample_d,
|
||||
grad_h_t_combined_d,
|
||||
@@ -2894,6 +2899,7 @@ impl IntegratedTrainer {
|
||||
&q_target_action_d,
|
||||
&self.sampled_rewards_d,
|
||||
&self.sampled_dones_d,
|
||||
&self.sampled_n_step_gammas_d,
|
||||
&self.isv_d,
|
||||
b_size,
|
||||
&mut target_dist_d,
|
||||
@@ -3495,6 +3501,7 @@ impl IntegratedTrainer {
|
||||
&q_target_action_d,
|
||||
&self.sampled_rewards_d,
|
||||
&self.sampled_dones_d,
|
||||
&self.sampled_n_step_gammas_d,
|
||||
&self.isv_d,
|
||||
b_size,
|
||||
&mut target_dist_d,
|
||||
@@ -4773,12 +4780,6 @@ impl IntegratedTrainer {
|
||||
fn push_to_replay(&mut self, b_size: usize) -> Result<()> {
|
||||
use crate::rl::common::{N_ACTIONS, Transition};
|
||||
|
||||
// Per-batch metadata via mapped-pinned staging (the only
|
||||
// permitted CPU↔GPU path per
|
||||
// `feedback_no_htod_htoh_only_mapped_pinned`). Each helper
|
||||
// allocates a mapped-pinned buffer, issues a DtoD copy from
|
||||
// the device source, syncs, and returns a host Vec read from
|
||||
// the page-mapped host_ptr.
|
||||
let actions_host = read_slice_i32_d(&self.stream, &self.actions_d, b_size)
|
||||
.context("push_to_replay: read actions_d")?;
|
||||
let rewards_host = read_slice_d(&self.stream, &self.rewards_d, b_size)
|
||||
@@ -4790,26 +4791,17 @@ impl IntegratedTrainer {
|
||||
let log_pi_old_host = read_slice_d(&self.stream, &self.log_pi_old_d, b_size)
|
||||
.context("push_to_replay: read log_pi_old_d")?;
|
||||
|
||||
// For each batch index, alloc per-transition device buffers
|
||||
// and DtoD-copy the per-batch slice of h_t / h_tp1 in.
|
||||
let n_step = self.isv_host[crate::rl::isv_slots::RL_N_STEP_INDEX] as usize;
|
||||
let n_step = n_step.max(1);
|
||||
let gamma = self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX];
|
||||
|
||||
let hidden_nbytes = HIDDEN_DIM * std::mem::size_of::<f32>();
|
||||
for b in 0..b_size {
|
||||
// Alloc h_t for this step's n-step entry.
|
||||
let mut h_t_per_b = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||||
.context("push_to_replay: alloc h_t per-batch")?;
|
||||
let mut h_tp1_per_b = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||||
.context("push_to_replay: alloc h_tp1 per-batch")?;
|
||||
|
||||
// Two separate inner scopes for the per-buffer DtoD so
|
||||
// each `SyncOnDrop` guard drops before the next
|
||||
// `device_ptr_mut` borrow / before the move into
|
||||
// `Transition`. Without this, the borrow checker rejects
|
||||
// the moves into `Transition { h_t, next_h_t, .. }` as
|
||||
// overlapping the guard lifetimes (canonical borrow-checker
|
||||
// dance for the cudarc raw-pointer API).
|
||||
let off = (b * HIDDEN_DIM * std::mem::size_of::<f32>()) as u64;
|
||||
unsafe {
|
||||
let s = self.stream.cu_stream();
|
||||
@@ -4820,6 +4812,50 @@ impl IntegratedTrainer {
|
||||
)
|
||||
.context("push_to_replay: DtoD per-batch h_t slice")?;
|
||||
}
|
||||
|
||||
let is_done = dones_host[b] > 0.5;
|
||||
self.n_step_buffer[b].push(NStepEntry {
|
||||
h_t: h_t_per_b,
|
||||
action: actions_host[b] as u32,
|
||||
raw_reward: raw_rewards_host[b],
|
||||
scaled_reward: rewards_host[b],
|
||||
done: is_done,
|
||||
log_pi_old: log_pi_old_host[b],
|
||||
});
|
||||
|
||||
// Flush when buffer reaches n OR a done event truncates.
|
||||
let should_flush = self.n_step_buffer[b].len() >= n_step || is_done;
|
||||
if !should_flush {
|
||||
continue;
|
||||
}
|
||||
|
||||
let buf = &self.n_step_buffer[b];
|
||||
let buf_len = buf.len();
|
||||
|
||||
// Compute n-step discounted return: R_n = Σ γᵏ rₖ
|
||||
let mut r_n_scaled = 0.0_f32;
|
||||
let mut r_n_raw = 0.0_f32;
|
||||
let mut gamma_acc = 1.0_f32;
|
||||
let mut any_done = false;
|
||||
for entry in buf.iter() {
|
||||
r_n_scaled += gamma_acc * entry.scaled_reward;
|
||||
r_n_raw += gamma_acc * entry.raw_reward;
|
||||
if entry.done {
|
||||
any_done = true;
|
||||
}
|
||||
gamma_acc *= gamma;
|
||||
}
|
||||
|
||||
// γⁿ for the bootstrap term. Zero if any done in window
|
||||
// (terminal state zeroes the bootstrap).
|
||||
let n_step_gamma_val = if any_done { 0.0 } else { gamma_acc };
|
||||
|
||||
// h_t from oldest entry, h_tp1 from current step.
|
||||
let oldest = &buf[0];
|
||||
let mut h_tp1_per_b = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||||
.context("push_to_replay: alloc h_tp1 per-batch")?;
|
||||
unsafe {
|
||||
let s = self.stream.cu_stream();
|
||||
let (h_tp1_base, _g2) = self.h_tp1_d.device_ptr(&self.stream);
|
||||
@@ -4829,15 +4865,36 @@ impl IntegratedTrainer {
|
||||
)
|
||||
.context("push_to_replay: DtoD per-batch h_tp1 slice")?;
|
||||
}
|
||||
|
||||
// Clone oldest h_t (it's about to be consumed by the drain).
|
||||
let mut h_t_oldest = self
|
||||
.stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM)
|
||||
.context("push_to_replay: alloc h_t_oldest")?;
|
||||
unsafe {
|
||||
let s = self.stream.cu_stream();
|
||||
let (src, _g1) = oldest.h_t.device_ptr(&self.stream);
|
||||
let (dst, _gd) = h_t_oldest.device_ptr_mut(&self.stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst, src, hidden_nbytes, s,
|
||||
)
|
||||
.context("push_to_replay: DtoD clone h_t_oldest")?;
|
||||
}
|
||||
let oldest_action = oldest.action;
|
||||
let oldest_log_pi = oldest.log_pi_old;
|
||||
|
||||
// Drain the buffer.
|
||||
self.n_step_buffer[b].clear();
|
||||
|
||||
self.replay.push(Transition {
|
||||
h_t: h_t_per_b,
|
||||
action: actions_host[b] as u32,
|
||||
reward: rewards_host[b],
|
||||
raw_reward: raw_rewards_host[b],
|
||||
h_t: h_t_oldest,
|
||||
action: oldest_action,
|
||||
reward: r_n_scaled,
|
||||
raw_reward: r_n_raw,
|
||||
next_h_t: h_tp1_per_b,
|
||||
done: dones_host[b] > 0.5,
|
||||
n_step_gamma: self.isv_host[crate::rl::isv_slots::RL_GAMMA_INDEX],
|
||||
log_pi_old: log_pi_old_host[b],
|
||||
done: any_done,
|
||||
n_step_gamma: n_step_gamma_val,
|
||||
log_pi_old: oldest_log_pi,
|
||||
q_value_old: [0.0; N_ACTIONS],
|
||||
});
|
||||
}
|
||||
@@ -4874,21 +4931,23 @@ impl IntegratedTrainer {
|
||||
let mut actions_host = vec![0i32; b_size];
|
||||
let mut rewards_host = vec![0.0f32; b_size];
|
||||
let mut dones_host = vec![0.0f32; b_size];
|
||||
let mut n_step_gammas_host = vec![0.0f32; b_size];
|
||||
let current_scale = self.isv_host[crate::rl::isv_slots::RL_REWARD_SCALE_INDEX];
|
||||
for (i, &idx) in indices.iter().enumerate() {
|
||||
let t = &self.replay.transitions[idx];
|
||||
actions_host[i] = t.action as i32;
|
||||
rewards_host[i] = t.raw_reward * current_scale;
|
||||
dones_host[i] = if t.done { 1.0 } else { 0.0 };
|
||||
n_step_gammas_host[i] = t.n_step_gamma;
|
||||
}
|
||||
// Mapped-pinned staging — the only permitted CPU→GPU path
|
||||
// per `feedback_no_htod_htoh_only_mapped_pinned`.
|
||||
write_slice_i32_d(&self.stream, &actions_host, &mut self.sampled_actions_d)
|
||||
.context("sample_and_gather: write sampled_actions_d")?;
|
||||
write_slice_f32_d(&self.stream, &rewards_host, &mut self.sampled_rewards_d)
|
||||
.context("sample_and_gather: write sampled_rewards_d")?;
|
||||
write_slice_f32_d(&self.stream, &dones_host, &mut self.sampled_dones_d)
|
||||
.context("sample_and_gather: write sampled_dones_d")?;
|
||||
write_slice_f32_d(&self.stream, &n_step_gammas_host, &mut self.sampled_n_step_gammas_d)
|
||||
.context("sample_and_gather: write sampled_n_step_gammas_d")?;
|
||||
|
||||
// Per-batch DtoD: per-transition h_t / h_tp1 device buffers →
|
||||
// contiguous `sampled_h_t_d` / `sampled_h_tp1_d` at slot i.
|
||||
|
||||
Reference in New Issue
Block a user