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

feat: add PPO CUDA kernels and curiosity inference kernel

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-31 00:39:59 +02:00
parent 56373f0941
commit 005fe591db
2 changed files with 77 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
crates/ml-ppo/src/*.cu Normal file
View File

77
crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu Normal file
View File

@@ -0,0 +1,77 @@
/**
* Curiosity inference-only kernel for Q-penalty computation.
*
* Runs the curiosity forward dynamics model (2-layer MLP) on a training batch
* and outputs per-sample MSE prediction error. This error is used by the C51
* loss kernel to scale down Bellman targets for novel/uncertain states.
*
* Architecture: [MARKET_DIM + 3] -> [CUR_HIDDEN] LeakyReLU -> [CUR_OUTPUT=MARKET_DIM]
* One thread per sample. Output: float prediction error per sample.
*
* Requires common_device_functions.cuh to be prepended (provides MARKET_DIM,
* CUR_INPUT, CUR_HIDDEN, CUR_OUTPUT, DQN_ORDER_ACTIONS, DQN_URGENCY_ACTIONS,
* DQN_NUM_ACTIONS, bf16 helpers).
*/
extern "C" __global__ void curiosity_inference_error(
const __nv_bfloat16* __restrict__ states, /* [N, state_dim] */
const int* __restrict__ actions, /* [N] */
const __nv_bfloat16* __restrict__ next_states, /* [N, state_dim] */
const __nv_bfloat16* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */
const __nv_bfloat16* __restrict__ b1, /* [CUR_HIDDEN] */
const __nv_bfloat16* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */
const __nv_bfloat16* __restrict__ b2, /* [CUR_OUTPUT] */
float* __restrict__ prediction_errors, /* [N] output: per-sample MSE */
int N,
int state_dim
) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= N) return;
const __nv_bfloat16* state_bf = states + tid * state_dim;
int action_idx = actions[tid];
const __nv_bfloat16* next_state_bf = next_states + tid * state_dim;
/* ── Build input: first MARKET_DIM state features + 3-class action one-hot ── */
__nv_bfloat16 input[CUR_INPUT];
for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; }
input[MARKET_DIM + 0] = bf16_zero();
input[MARKET_DIM + 1] = bf16_zero();
input[MARKET_DIM + 2] = bf16_zero();
/* Action to category one-hot — decode exposure from factored action.
* Factored action = exposure_idx * (b1_size * b2_size) + order_idx * b2_size + urgency_idx.
* With b0=9, b1=3, b2=3: exposure_idx = action / 9. */
int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS);
int category;
if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */
else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */
else category = 2; /* Long */
input[MARKET_DIM + category] = bf16_one();
/* ── Layer 1: hidden = LeakyReLU(w1 * input + b1, alpha=0.01) ── */
float hidden[CUR_HIDDEN];
for (int h = 0; h < CUR_HIDDEN; h++) {
float sum = (float)b1[h];
for (int i = 0; i < CUR_INPUT; i++) {
sum += (float)w1[h * CUR_INPUT + i] * (float)input[i];
}
hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
}
/* ── Layer 2: pred = w2 * hidden + b2 (no activation) ── */
float mse = 0.0f;
for (int o = 0; o < CUR_OUTPUT; o++) {
float sum = (float)b2[o];
for (int h = 0; h < CUR_HIDDEN; h++) {
sum += (float)w2[o * CUR_HIDDEN + h] * hidden[h];
}
/* Accumulate MSE against actual next state */
float diff = sum - (float)next_state_bf[o];
mse += diff * diff;
}
mse /= (float)CUR_OUTPUT;
/* Clamp to prevent extreme penalties from corrupting training */
prediction_errors[tid] = fminf(mse, 10.0f);
}