13bf277cd6
Implements §4.2/4.3/4.4 + §5 of Phase 4 spec.
Three new CUDA kernels:
1. rl_dueling_q_bellman_target.cu — argmax over target composed_Q at
s_{t+1}, build target_value = r + γ^n × (1-done) × max_Q. Grid
(B,1,1), block (N_ACTIONS,1,1). Reads γ via per-sample n_step_gammas
passed by trainer (matches PER convention).
2. rl_dueling_q_loss_and_grad.cu — scalar Huber loss on
(target − online_composed_Q[taken]). Emits per-batch loss and
grad_composed (only taken action has nonzero gradient — this is
single-Q scalar regression, not distributional CE). Grid (B,1,1),
block (N_ACTIONS,1,1). Threads write zero into non-taken cells.
3. rl_dueling_q_decompose_and_bwd.cu — decompose grad_composed →
grad_V + grad_A via mean-subtraction Jacobian:
grad_V[b] = Σ_a grad_composed[b, a] = grad_composed[b, a_taken]
grad_A[b, a] = grad_composed[b, a] − (1/N) × grad_V[b]
Then computes per-batch weight gradients via outer product with h_t:
grad_w_v_pb[b, c] = grad_V[b] × h_t[b, c]
grad_b_v_pb[b] = grad_V[b]
grad_w_a_pb[b, c, a] = grad_A[b, a] × h_t[b, c]
grad_b_a_pb[b, a] = grad_A[b, a]
Grid (B,1,1), block (HIDDEN_DIM,1,1). Thread 0 also writes biases.
DuelingQHead Rust API (3 new methods):
- build_bellman_target(target_composed_q, r, dones, n_step_gammas, B, out)
- compute_loss_and_grad(online_composed_q, target_value, actions, B,
loss_pb, grad_composed_out)
- decompose_and_backward_to_weights(h_t, grad_composed, actions, B,
grad_w_v_pb, grad_b_v_pb,
grad_w_a_pb, grad_b_a_pb)
Caller (trainer) is responsible for reduce_axis0 of per-batch grads
→ final weight gradients [HIDDEN_DIM], [1], [HIDDEN_DIM × N_ACTIONS],
[N_ACTIONS]. Will be wired in Phase 4.2.
Soft-update of target weights still TODO — recommend reusing
dqn_target_soft_update kernel pattern in Phase 4.2.
Validated:
- cargo build --release clean
- integrated_trainer_smoke (1 step) passes
Per pearl_complement_internal_loss_vs_external_consumers: this loss
path is fully internal to DuelingQHead. No downstream consumer ever
sees composed_Q or grad_composed. The four prior session failure
modes are structurally impossible.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
90 lines
3.3 KiB
Plaintext
90 lines
3.3 KiB
Plaintext
// rl_dueling_q_decompose_and_bwd.cu — Phase 4 decompose grad + weight grad.
|
||
//
|
||
// Decompose:
|
||
// composed_Q[b, a] = V[b] + A[b, a] − (1/N) Σ_a' A[b, a']
|
||
//
|
||
// Chain rule (only taken action has nonzero grad_composed[b, a]):
|
||
// grad_V[b] = Σ_a grad_composed[b, a] = grad_composed[b, a_taken]
|
||
// grad_A[b, a] = grad_composed[b, a] − (1/N) × grad_V[b]
|
||
//
|
||
// For a == a_taken: grad_A[b, a] = (1 − 1/N) × grad_composed[b, a_taken]
|
||
// For a ≠ a_taken: grad_A[b, a] = −(1/N) × grad_composed[b, a_taken]
|
||
//
|
||
// After decompose, compute per-batch weight gradients via outer
|
||
// product with h_t:
|
||
// grad_w_v_pb[b, c] = grad_V[b] × h_t[b, c]
|
||
// grad_b_v_pb[b] = grad_V[b]
|
||
// grad_w_a_pb[b, c, a] = grad_A[b, a] × h_t[b, c]
|
||
// grad_b_a_pb[b, a] = grad_A[b, a]
|
||
//
|
||
// Caller reduces per-batch grads via reduce_axis0 to final shapes:
|
||
// grad_w_v [HIDDEN_DIM]
|
||
// grad_b_v [1]
|
||
// grad_w_a [HIDDEN_DIM × N_ACTIONS]
|
||
// grad_b_a [N_ACTIONS]
|
||
//
|
||
// Block layout:
|
||
// grid = (B, 1, 1)
|
||
// block = (HIDDEN_DIM, 1, 1) — one thread per hidden-dim index
|
||
// Each thread loops over N_ACTIONS to write A weights, plus the
|
||
// single V weight. Thread 0 additionally writes grad_b_v_pb +
|
||
// grad_b_a_pb (small).
|
||
//
|
||
// Per feedback_no_atomicadd: sole-writer per (b, c, a) and (b, c) cells.
|
||
|
||
#define HIDDEN_DIM 128
|
||
#define N_ACTIONS 11
|
||
|
||
extern "C" __global__ void rl_dueling_q_decompose_and_weight_grad(
|
||
const float* __restrict__ h_t, // [B × HIDDEN_DIM]
|
||
const float* __restrict__ grad_composed, // [B × N_ACTIONS] (only taken cell nonzero)
|
||
const int* __restrict__ actions_taken, // [B]
|
||
int B,
|
||
float* __restrict__ grad_w_v_pb, // [B × HIDDEN_DIM]
|
||
float* __restrict__ grad_b_v_pb, // [B]
|
||
float* __restrict__ grad_w_a_pb, // [B × HIDDEN_DIM × N_ACTIONS]
|
||
float* __restrict__ grad_b_a_pb // [B × N_ACTIONS]
|
||
) {
|
||
const int b = blockIdx.x;
|
||
const int c = threadIdx.x;
|
||
if (b >= B) return;
|
||
if (c >= HIDDEN_DIM) return;
|
||
|
||
int a_t = actions_taken[b];
|
||
if (a_t < 0) a_t = 0;
|
||
if (a_t >= N_ACTIONS) a_t = 0;
|
||
|
||
// grad_composed is nonzero only at a == a_t.
|
||
const float gc_taken = grad_composed[b * N_ACTIONS + a_t];
|
||
|
||
// grad_V[b] = Σ_a grad_composed[b, a] = gc_taken (others are 0).
|
||
const float grad_v = gc_taken;
|
||
|
||
// h_t[b, c].
|
||
const float h_bc = h_t[b * HIDDEN_DIM + c];
|
||
|
||
// Per-batch V weight grad.
|
||
grad_w_v_pb[b * HIDDEN_DIM + c] = grad_v * h_bc;
|
||
|
||
// Per-batch A weight grad (one thread writes N_ACTIONS values).
|
||
// grad_A[b, a] = grad_composed[b, a] − (1/N) × grad_V[b]
|
||
// = (a == a_t ? gc_taken : 0) − (1/N) × gc_taken
|
||
const float inv_N = 1.0f / (float)N_ACTIONS;
|
||
#pragma unroll
|
||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||
const float grad_a = (a == a_t ? gc_taken : 0.0f) - inv_N * gc_taken;
|
||
// Layout: grad_w_a_pb[b, c, a] = h_bc * grad_a
|
||
grad_w_a_pb[b * HIDDEN_DIM * N_ACTIONS + c * N_ACTIONS + a] = h_bc * grad_a;
|
||
}
|
||
|
||
// Thread 0 writes biases (small).
|
||
if (c == 0) {
|
||
grad_b_v_pb[b] = grad_v;
|
||
#pragma unroll
|
||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||
const float grad_a = (a == a_t ? gc_taken : 0.0f) - inv_N * gc_taken;
|
||
grad_b_a_pb[b * N_ACTIONS + a] = grad_a;
|
||
}
|
||
}
|
||
}
|