feat(cuda): add PPO experience kernel — actor, critic, GAE, full rollout
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -477,4 +477,28 @@ mod tests {
|
||||
assert!((cfg.barrier_profit_mult - 1.02).abs() < f32::EPSILON);
|
||||
assert!((cfg.barrier_loss_mult - 0.98).abs() < f32::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ppo_kernel_source_contains_all_functions() {
|
||||
let src = include_str!("ppo_experience_kernel.cu");
|
||||
assert!(src.contains("ppo_actor_forward"), "Missing ppo_actor_forward");
|
||||
assert!(src.contains("softmax_sample"), "Missing softmax_sample");
|
||||
assert!(src.contains("ppo_critic_forward"), "Missing ppo_critic_forward");
|
||||
assert!(src.contains("compute_gae_backward"), "Missing compute_gae_backward");
|
||||
assert!(src.contains("ppo_full_experience_kernel"), "Missing kernel entry point");
|
||||
assert!(src.contains("extern \"C\""), "Missing extern C linkage");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ppo_kernel_source_concatenation() {
|
||||
let common = include_str!("common_device_functions.cuh");
|
||||
let kernel = include_str!("ppo_experience_kernel.cu");
|
||||
let full = format!("{}\n{}", common, kernel);
|
||||
assert!(full.contains("gpu_random"));
|
||||
assert!(full.contains("matvec_leaky_relu"));
|
||||
assert!(full.contains("ppo_full_experience_kernel"));
|
||||
assert!(full.contains("compute_gae_backward"));
|
||||
assert!(full.contains("softmax_sample"));
|
||||
assert!(!full.contains("q_forward_dueling"));
|
||||
}
|
||||
}
|
||||
|
||||
677
crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu
Normal file
677
crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu
Normal file
@@ -0,0 +1,677 @@
|
||||
/**
|
||||
* Zero-Roundtrip PPO Experience Collection Kernel
|
||||
*
|
||||
* Requires common_device_functions.cuh prepended via NVRTC source concatenation.
|
||||
* Launch config: grid=(ceil(N/32),1,1), block=(32,1,1).
|
||||
* Each thread runs one independent episode.
|
||||
*
|
||||
* Phase A: Forward rollout (L timesteps) — actor, critic, portfolio, rewards
|
||||
* Phase B: Backward GAE scan — advantages and returns
|
||||
*/
|
||||
|
||||
/* PPO Actor layer sizes: 54 -> 128 (ReLU) -> 64 (ReLU) -> 45 */
|
||||
#define ACTOR_H1 128
|
||||
#define ACTOR_H2 64
|
||||
|
||||
/* PPO Critic layer sizes: 54 -> 512 -> 384 -> 256 -> 128 -> 64 -> 1 (all ReLU) */
|
||||
#define CRITIC_H1 512
|
||||
#define CRITIC_H2 384
|
||||
#define CRITIC_H3 256
|
||||
#define CRITIC_H4 128
|
||||
#define CRITIC_H5 64
|
||||
|
||||
/* Maximum GAE rollout length (matches timesteps_per_episode default) */
|
||||
#define MAX_GAE_LEN 500
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* PPO-Specific Device Functions */
|
||||
/* ------------------------------------------------------------------ */
|
||||
|
||||
/**
|
||||
* PPO Actor MLP forward pass.
|
||||
*
|
||||
* state[54] -> h1[128] (LeakyReLU) -> h2[64] (LeakyReLU) -> logits[45]
|
||||
*
|
||||
* @param state Input state vector [STATE_DIM]
|
||||
* @param pw1 Layer 1 weights [ACTOR_H1, STATE_DIM]
|
||||
* @param pb1 Layer 1 biases [ACTOR_H1]
|
||||
* @param pw2 Layer 2 weights [ACTOR_H2, ACTOR_H1]
|
||||
* @param pb2 Layer 2 biases [ACTOR_H2]
|
||||
* @param pw3 Output weights [NUM_ACTIONS, ACTOR_H2]
|
||||
* @param pb3 Output biases [NUM_ACTIONS]
|
||||
* @param h1 Scratch buffer [ACTOR_H1]
|
||||
* @param h2 Scratch buffer [ACTOR_H2]
|
||||
* @param logits Output logits [NUM_ACTIONS]
|
||||
*/
|
||||
__device__ void ppo_actor_forward(
|
||||
const float* state,
|
||||
const float* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] */
|
||||
const float* __restrict__ pb1, /* [ACTOR_H1] */
|
||||
const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] */
|
||||
const float* __restrict__ pb2, /* [ACTOR_H2] */
|
||||
const float* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] */
|
||||
const float* __restrict__ pb3, /* [NUM_ACTIONS] */
|
||||
float* h1, /* [ACTOR_H1] scratch */
|
||||
float* h2, /* [ACTOR_H2] scratch */
|
||||
float* logits /* [NUM_ACTIONS] output */
|
||||
) {
|
||||
/* Hidden layer 1: state -> h1 with LeakyReLU */
|
||||
matvec_leaky_relu(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);
|
||||
|
||||
/* Hidden layer 2: h1 -> h2 with LeakyReLU */
|
||||
matvec_leaky_relu(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
|
||||
|
||||
/* Output layer: h2 -> logits (no activation) */
|
||||
matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Stable softmax + categorical sampling.
|
||||
*
|
||||
* 1. Find max logit for numerical stability
|
||||
* 2. exp(logit - max) and accumulate sum
|
||||
* 3. Normalize to probabilities
|
||||
* 4. CDF scan with LCG random draw -> action index
|
||||
* 5. Compute log(p[action]) for PPO loss
|
||||
*
|
||||
* @param logits Input logits [NUM_ACTIONS]
|
||||
* @param probs Scratch + output probabilities [NUM_ACTIONS]
|
||||
* @param rng Pointer to LCG RNG state
|
||||
* @param out_action Output: selected action index
|
||||
* @param out_log_prob Output: log probability of selected action
|
||||
*/
|
||||
__device__ void softmax_sample(
|
||||
const float* logits,
|
||||
float* probs,
|
||||
unsigned int* rng,
|
||||
int* out_action,
|
||||
float* out_log_prob
|
||||
) {
|
||||
/* Step 1: Find max logit for numerical stability */
|
||||
float max_logit = logits[0];
|
||||
for (int i = 1; i < NUM_ACTIONS; i++) {
|
||||
if (logits[i] > max_logit) max_logit = logits[i];
|
||||
}
|
||||
|
||||
/* Step 2: exp(logit - max) and accumulate sum */
|
||||
float sum_exp = 0.0f;
|
||||
for (int i = 0; i < NUM_ACTIONS; i++) {
|
||||
probs[i] = expf(logits[i] - max_logit);
|
||||
sum_exp += probs[i];
|
||||
}
|
||||
|
||||
/* Step 3: Normalize to probabilities */
|
||||
float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
|
||||
for (int i = 0; i < NUM_ACTIONS; i++) {
|
||||
probs[i] *= inv_sum;
|
||||
}
|
||||
|
||||
/* Step 4: CDF scan + random draw -> action index */
|
||||
float u = gpu_random(rng);
|
||||
float cdf = 0.0f;
|
||||
int action = NUM_ACTIONS - 1; /* default to last action */
|
||||
for (int i = 0; i < NUM_ACTIONS; i++) {
|
||||
cdf += probs[i];
|
||||
if (u < cdf) {
|
||||
action = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
/* Step 5: log probability for PPO loss */
|
||||
float p = fmaxf(probs[action], 1e-8f); /* clamp for log safety */
|
||||
*out_action = action;
|
||||
*out_log_prob = logf(p);
|
||||
}
|
||||
|
||||
/**
|
||||
* PPO Critic (value network) forward pass — 5-layer deep with ping-pong buffers.
|
||||
*
|
||||
* state[54] -> 512(LReLU) -> 384(LReLU) -> 256(LReLU) -> 128(LReLU) -> 64(LReLU) -> 1
|
||||
*
|
||||
* Uses two scratch buffers (scratch_a[512], scratch_b[512]) that alternate
|
||||
* between layers to avoid extra memory.
|
||||
*
|
||||
* @param state Input state vector [STATE_DIM]
|
||||
* @param vw1-vw6 Weight matrices for each layer
|
||||
* @param vb1-vb6 Bias vectors for each layer
|
||||
* @param scratch_a Ping-pong scratch buffer A [CRITIC_H1] (512 wide)
|
||||
* @param scratch_b Ping-pong scratch buffer B [CRITIC_H1] (512 wide)
|
||||
* @return Scalar value estimate V(s)
|
||||
*/
|
||||
__device__ float ppo_critic_forward(
|
||||
const float* state,
|
||||
const float* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */
|
||||
const float* __restrict__ vb1, /* [CRITIC_H1] = [512] */
|
||||
const float* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */
|
||||
const float* __restrict__ vb2, /* [CRITIC_H2] = [384] */
|
||||
const float* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */
|
||||
const float* __restrict__ vb3, /* [CRITIC_H3] = [256] */
|
||||
const float* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */
|
||||
const float* __restrict__ vb4, /* [CRITIC_H4] = [128] */
|
||||
const float* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */
|
||||
const float* __restrict__ vb5, /* [CRITIC_H5] = [64] */
|
||||
const float* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */
|
||||
const float* __restrict__ vb6, /* [1] */
|
||||
float* scratch_a, /* [CRITIC_H1] = [512] ping-pong A */
|
||||
float* scratch_b /* [CRITIC_H1] = [512] ping-pong B */
|
||||
) {
|
||||
/* Layer 1: state[54] -> scratch_a[512] with LeakyReLU */
|
||||
matvec_leaky_relu(vw1, vb1, state, scratch_a, STATE_DIM, CRITIC_H1, 1);
|
||||
|
||||
/* Layer 2: scratch_a[512] -> scratch_b[384] with LeakyReLU */
|
||||
matvec_leaky_relu(vw2, vb2, scratch_a, scratch_b, CRITIC_H1, CRITIC_H2, 1);
|
||||
|
||||
/* Layer 3: scratch_b[384] -> scratch_a[256] with LeakyReLU */
|
||||
matvec_leaky_relu(vw3, vb3, scratch_b, scratch_a, CRITIC_H2, CRITIC_H3, 1);
|
||||
|
||||
/* Layer 4: scratch_a[256] -> scratch_b[128] with LeakyReLU */
|
||||
matvec_leaky_relu(vw4, vb4, scratch_a, scratch_b, CRITIC_H3, CRITIC_H4, 1);
|
||||
|
||||
/* Layer 5: scratch_b[128] -> scratch_a[64] with LeakyReLU */
|
||||
matvec_leaky_relu(vw5, vb5, scratch_b, scratch_a, CRITIC_H4, CRITIC_H5, 1);
|
||||
|
||||
/* Output layer: scratch_a[64] -> scalar (no activation) */
|
||||
float value = vb6[0];
|
||||
for (int i = 0; i < CRITIC_H5; i++) {
|
||||
value += vw6[i] * scratch_a[i];
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generalized Advantage Estimation (GAE) backward scan.
|
||||
*
|
||||
* Computes advantages and returns by scanning backwards through
|
||||
* the collected rollout data:
|
||||
*
|
||||
* for t = L-1 down to 0:
|
||||
* delta = rewards[t] + gamma * values[t+1] * (1-dones[t]) - values[t]
|
||||
* gae = delta + gamma * lambda * (1-dones[t]) * gae
|
||||
* advantages[t] = gae
|
||||
* returns[t] = gae + values[t]
|
||||
*
|
||||
* @param rewards Per-timestep rewards [L]
|
||||
* @param values Per-timestep value estimates [L+1] (values[L] is bootstrap)
|
||||
* @param dones Per-timestep done flags [L] (1.0 = done, 0.0 = not done)
|
||||
* @param advantages Output advantage estimates [L]
|
||||
* @param returns Output return targets [L]
|
||||
* @param L Number of timesteps
|
||||
* @param gamma Discount factor
|
||||
* @param lambda GAE lambda parameter
|
||||
*/
|
||||
__device__ void compute_gae_backward(
|
||||
const float* rewards,
|
||||
const float* values,
|
||||
const float* dones,
|
||||
float* advantages,
|
||||
float* returns,
|
||||
int L,
|
||||
float gamma,
|
||||
float lambda
|
||||
) {
|
||||
float gae = 0.0f;
|
||||
|
||||
for (int t = L - 1; t >= 0; t--) {
|
||||
float not_done = 1.0f - dones[t];
|
||||
float delta = rewards[t] + gamma * values[t + 1] * not_done - values[t];
|
||||
gae = delta + gamma * lambda * not_done * gae;
|
||||
advantages[t] = gae;
|
||||
returns[t] = gae + values[t];
|
||||
}
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* Main Kernel */
|
||||
/* ------------------------------------------------------------------ */
|
||||
|
||||
/**
|
||||
* Full PPO experience collection kernel.
|
||||
*
|
||||
* Each thread runs one independent episode of L timesteps (Phase A),
|
||||
* then performs a backward GAE scan (Phase B).
|
||||
*
|
||||
* Grid: (ceil(N/32), 1, 1), Block: (32, 1, 1).
|
||||
*/
|
||||
extern "C" __global__ void ppo_full_experience_kernel(
|
||||
/* ---- Market data ---- */
|
||||
const float* __restrict__ market_features, /* [total_bars, MARKET_DIM] */
|
||||
const float* __restrict__ targets, /* [total_bars, 4] */
|
||||
const int* __restrict__ episode_starts, /* [N] */
|
||||
|
||||
/* ---- Actor weights (6 pointers) ---- */
|
||||
const float* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] = [128, 54] */
|
||||
const float* __restrict__ pb1, /* [ACTOR_H1] = [128] */
|
||||
const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */
|
||||
const float* __restrict__ pb2, /* [ACTOR_H2] = [64] */
|
||||
const float* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] = [45, 64] */
|
||||
const float* __restrict__ pb3, /* [NUM_ACTIONS] = [45] */
|
||||
|
||||
/* ---- Critic weights (12 pointers) ---- */
|
||||
const float* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */
|
||||
const float* __restrict__ vb1, /* [CRITIC_H1] = [512] */
|
||||
const float* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */
|
||||
const float* __restrict__ vb2, /* [CRITIC_H2] = [384] */
|
||||
const float* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */
|
||||
const float* __restrict__ vb3, /* [CRITIC_H3] = [256] */
|
||||
const float* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */
|
||||
const float* __restrict__ vb4, /* [CRITIC_H4] = [128] */
|
||||
const float* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */
|
||||
const float* __restrict__ vb5, /* [CRITIC_H5] = [64] */
|
||||
const float* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */
|
||||
const float* __restrict__ vb6, /* [1] */
|
||||
|
||||
/* ---- Curiosity model weights (4 pointers) ---- */
|
||||
const float* __restrict__ cur_w1, /* [CUR_HIDDEN, CUR_INPUT] */
|
||||
const float* __restrict__ cur_b1, /* [CUR_HIDDEN] */
|
||||
const float* __restrict__ cur_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */
|
||||
const float* __restrict__ cur_b2, /* [CUR_OUTPUT] */
|
||||
|
||||
/* ---- Per-episode mutable state arrays ---- */
|
||||
float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */
|
||||
float* barrier_states, /* [N, BARRIER_STATE_SIZE] */
|
||||
int* diversity_windows, /* [N, DIVERSITY_WINDOW] */
|
||||
int* diversity_metas, /* [N, 2] */
|
||||
|
||||
/* ---- Barrier config (shared) ---- */
|
||||
const float* __restrict__ barrier_config, /* [3]: profit_mult, loss_mult, max_bars */
|
||||
|
||||
/* ---- Scalar configs ---- */
|
||||
float max_position,
|
||||
int episode_length,
|
||||
int total_bars,
|
||||
int L, /* timesteps per episode (rollout length) */
|
||||
float gamma,
|
||||
float gae_lambda,
|
||||
float curiosity_max_reward,
|
||||
int N, /* total number of episodes */
|
||||
float barrier_scale,
|
||||
float diversity_scale,
|
||||
float curiosity_scale,
|
||||
float risk_weight,
|
||||
|
||||
/* ---- RNG states [N] ---- */
|
||||
unsigned int* rng_states,
|
||||
|
||||
/* ---- Output arrays ---- */
|
||||
float* out_states, /* [N, L, STATE_DIM] */
|
||||
int* out_actions, /* [N, L] */
|
||||
float* out_log_probs, /* [N, L] */
|
||||
float* out_advantages, /* [N, L] */
|
||||
float* out_returns, /* [N, L] */
|
||||
int* out_dones /* [N, L] */
|
||||
) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (tid >= N) return;
|
||||
|
||||
/* ---- Load per-thread portfolio state ---- */
|
||||
int ps_off = tid * PORTFOLIO_STATE_SIZE;
|
||||
float cash = portfolio_states[ps_off + 0];
|
||||
float position = portfolio_states[ps_off + 1];
|
||||
float entry_price = portfolio_states[ps_off + 2];
|
||||
float initial_cap = portfolio_states[ps_off + 3];
|
||||
float spread = portfolio_states[ps_off + 4];
|
||||
float last_price = portfolio_states[ps_off + 5];
|
||||
float reserve_pct = portfolio_states[ps_off + 6];
|
||||
float cum_costs = portfolio_states[ps_off + 7];
|
||||
|
||||
/* ---- Load per-thread barrier state ---- */
|
||||
int bs_off = tid * BARRIER_STATE_SIZE;
|
||||
float barrier_st[BARRIER_STATE_SIZE];
|
||||
for (int i = 0; i < BARRIER_STATE_SIZE; i++)
|
||||
barrier_st[i] = barrier_states[bs_off + i];
|
||||
|
||||
/* ---- Load per-thread diversity window ---- */
|
||||
int dw_off = tid * DIVERSITY_WINDOW;
|
||||
int div_window[DIVERSITY_WINDOW];
|
||||
for (int i = 0; i < DIVERSITY_WINDOW; i++)
|
||||
div_window[i] = diversity_windows[dw_off + i];
|
||||
|
||||
int dm_off = tid * 2;
|
||||
int div_meta[2];
|
||||
div_meta[0] = diversity_metas[dm_off + 0];
|
||||
div_meta[1] = diversity_metas[dm_off + 1];
|
||||
|
||||
unsigned int rng = rng_states[tid];
|
||||
int ep_start = episode_starts[tid];
|
||||
|
||||
/* ---- Per-thread scratch buffers ---- */
|
||||
float state[STATE_DIM];
|
||||
float actor_h1[ACTOR_H1];
|
||||
float actor_h2[ACTOR_H2];
|
||||
float logits[NUM_ACTIONS];
|
||||
float probs[NUM_ACTIONS];
|
||||
float critic_a[CRITIC_H1]; /* ping-pong buffer A (512 wide) */
|
||||
float critic_b[CRITIC_H1]; /* ping-pong buffer B (512 wide) */
|
||||
float next_state[STATE_DIM];
|
||||
float cur_scratch[CUR_HIDDEN];
|
||||
|
||||
/* GAE accumulation arrays */
|
||||
float gae_values[MAX_GAE_LEN + 1]; /* values[L] is bootstrap */
|
||||
float gae_rewards[MAX_GAE_LEN];
|
||||
float gae_dones[MAX_GAE_LEN];
|
||||
float gae_advantages[MAX_GAE_LEN];
|
||||
float gae_returns[MAX_GAE_LEN];
|
||||
|
||||
int step_in_episode = 0;
|
||||
int actual_L = (L <= MAX_GAE_LEN) ? L : MAX_GAE_LEN;
|
||||
|
||||
/* ================================================================ */
|
||||
/* Phase A: Forward Rollout (L timesteps) */
|
||||
/* ================================================================ */
|
||||
|
||||
for (int t = 0; t < actual_L; t++) {
|
||||
int global_bar = ep_start + t;
|
||||
int out_off = tid * actual_L + t;
|
||||
|
||||
/* Handle out-of-data */
|
||||
if (global_bar >= total_bars - 1) {
|
||||
for (int i = 0; i < STATE_DIM; i++)
|
||||
out_states[out_off * STATE_DIM + i] = 0.0f;
|
||||
out_actions[out_off] = 0;
|
||||
out_log_probs[out_off] = 0.0f;
|
||||
out_dones[out_off] = 1;
|
||||
gae_values[t] = 0.0f;
|
||||
gae_rewards[t] = 0.0f;
|
||||
gae_dones[t] = 1.0f;
|
||||
continue;
|
||||
}
|
||||
|
||||
/* ---- Step 1: Read 51 market features from global memory ---- */
|
||||
int mf_off = global_bar * MARKET_DIM;
|
||||
for (int i = 0; i < MARKET_DIM; i++)
|
||||
state[i] = market_features[mf_off + i];
|
||||
|
||||
/* ---- Step 2: Compute 3 portfolio features ---- */
|
||||
int t_off = global_bar * 4;
|
||||
float current_close = targets[t_off + 0];
|
||||
float next_close = targets[t_off + 1];
|
||||
float current_close_raw = targets[t_off + 2];
|
||||
float next_close_raw = targets[t_off + 3];
|
||||
|
||||
float price = (current_close_raw != 0.0f) ? current_close_raw : current_close;
|
||||
if (price <= 0.0f) price = 1.0f;
|
||||
|
||||
float current_value = cash + position * price;
|
||||
float current_norm = current_value / initial_cap;
|
||||
float max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f;
|
||||
float pos_norm = position / max_pos_norm;
|
||||
|
||||
state[MARKET_DIM + 0] = current_norm; /* normalized value */
|
||||
state[MARKET_DIM + 1] = pos_norm; /* normalized position */
|
||||
state[MARKET_DIM + 2] = spread; /* spread */
|
||||
|
||||
/* ---- Step 3: Actor forward -> logits -> softmax_sample ---- */
|
||||
ppo_actor_forward(
|
||||
state,
|
||||
pw1, pb1, pw2, pb2, pw3, pb3,
|
||||
actor_h1, actor_h2, logits
|
||||
);
|
||||
|
||||
int action_idx;
|
||||
float log_prob;
|
||||
softmax_sample(logits, probs, &rng, &action_idx, &log_prob);
|
||||
|
||||
/* ---- Step 4: Critic forward -> value estimate ---- */
|
||||
float value = ppo_critic_forward(
|
||||
state,
|
||||
vw1, vb1, vw2, vb2, vw3, vb3,
|
||||
vw4, vb4, vw5, vb5, vw6, vb6,
|
||||
critic_a, critic_b
|
||||
);
|
||||
|
||||
/* Store value for GAE */
|
||||
gae_values[t] = value;
|
||||
|
||||
/* ---- Step 5: Portfolio simulation ---- */
|
||||
float target_exposure = action_to_exposure(action_idx);
|
||||
float target_position = target_exposure * max_position;
|
||||
float tx_rate = action_to_tx_cost(action_idx);
|
||||
|
||||
/* Detect reversal (sign change) */
|
||||
int is_reversal = (position > 0.0f && target_position < 0.0f) ||
|
||||
(position < 0.0f && target_position > 0.0f);
|
||||
|
||||
if (is_reversal) {
|
||||
/* Phase 1: Close current position */
|
||||
float close_cash = position * price;
|
||||
float close_cost = fabsf(position) * price * tx_rate;
|
||||
cash += close_cash - close_cost;
|
||||
cum_costs += close_cost;
|
||||
|
||||
/* Phase 2: Open opposite position */
|
||||
float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f;
|
||||
float affordable = fmaxf(cash - reserve, 0.0f);
|
||||
float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f;
|
||||
max_contracts = floorf(max_contracts);
|
||||
float actual = fminf(max_contracts, fabsf(target_position));
|
||||
|
||||
if (actual > 0.0f) {
|
||||
float new_pos = (target_position > 0.0f) ? actual : -actual;
|
||||
float open_cost = actual * price * tx_rate;
|
||||
cash -= new_pos * price + open_cost;
|
||||
cum_costs += open_cost;
|
||||
position = new_pos;
|
||||
entry_price = price;
|
||||
} else {
|
||||
position = 0.0f;
|
||||
entry_price = 0.0f;
|
||||
}
|
||||
} else {
|
||||
/* Non-reversal: adjust position directly */
|
||||
float delta = target_position - position;
|
||||
if (fabsf(delta) > 0.0f) {
|
||||
float trade_cost = fabsf(delta) * price * tx_rate;
|
||||
cum_costs += trade_cost;
|
||||
cash -= trade_cost;
|
||||
|
||||
/* Cash reserve check for buys */
|
||||
if (delta > 0.0f && reserve_pct > 0.0f) {
|
||||
float pv = cash + position * price;
|
||||
float reserve = pv * (reserve_pct / 100.0f);
|
||||
float buy_cost = delta * price;
|
||||
if (cash - buy_cost < reserve) {
|
||||
float affordable = fmaxf(cash - reserve, 0.0f);
|
||||
delta = fminf(delta, (price > 0.0f) ? floorf(affordable / price) : 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
if (delta > 0.0f) {
|
||||
entry_price = price;
|
||||
} else if (target_position == 0.0f) {
|
||||
entry_price = 0.0f;
|
||||
}
|
||||
cash -= delta * price;
|
||||
position = position + delta;
|
||||
}
|
||||
}
|
||||
last_price = price;
|
||||
|
||||
/* ---- Step 6: Barrier tracking ---- */
|
||||
float old_barrier_entry = barrier_st[0];
|
||||
if (entry_price > 0.0f && old_barrier_entry <= 0.0f) {
|
||||
barrier_init(barrier_st, barrier_config, entry_price, global_bar);
|
||||
}
|
||||
|
||||
int barrier_label = barrier_check(barrier_st, price, global_bar, position);
|
||||
if (barrier_label != 0) {
|
||||
barrier_reset(barrier_st);
|
||||
}
|
||||
|
||||
/* ---- Step 7: Diversity entropy penalty ---- */
|
||||
float div_penalty = diversity_entropy(div_window, div_meta, action_idx);
|
||||
|
||||
/* ---- Step 8: Mark-to-market + build next_state for curiosity ---- */
|
||||
float next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close;
|
||||
if (next_price <= 0.0f) next_price = price;
|
||||
float next_value = cash + position * next_price;
|
||||
float next_norm = next_value / initial_cap;
|
||||
|
||||
int next_bar = global_bar + 1;
|
||||
if (next_bar < total_bars) {
|
||||
int nmf_off = next_bar * MARKET_DIM;
|
||||
for (int i = 0; i < MARKET_DIM; i++)
|
||||
next_state[i] = market_features[nmf_off + i];
|
||||
} else {
|
||||
for (int i = 0; i < MARKET_DIM; i++)
|
||||
next_state[i] = state[i];
|
||||
}
|
||||
float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f;
|
||||
next_state[MARKET_DIM + 0] = next_norm;
|
||||
next_state[MARKET_DIM + 1] = position / next_max_pos_norm;
|
||||
next_state[MARKET_DIM + 2] = spread;
|
||||
|
||||
/* ---- Step 9: Curiosity inference ---- */
|
||||
float curiosity_reward = 0.0f;
|
||||
if (cur_w1 != 0) {
|
||||
curiosity_reward = curiosity_inference(
|
||||
state, next_state, action_idx,
|
||||
cur_w1, cur_b1, cur_w2, cur_b2,
|
||||
cur_scratch, curiosity_max_reward
|
||||
);
|
||||
}
|
||||
|
||||
/* ---- Step 10: Risk penalty (drawdown) ---- */
|
||||
float pnl_reward = 0.0f;
|
||||
if (current_norm > 0.0f) {
|
||||
pnl_reward = (next_norm - current_norm) / current_norm;
|
||||
}
|
||||
|
||||
float abs_pos = fabsf(pos_norm);
|
||||
if (abs_pos > 0.8f) {
|
||||
pnl_reward -= (abs_pos - 0.8f) * 5.0f * risk_weight;
|
||||
}
|
||||
|
||||
/* ---- Step 11: Reward combination ---- */
|
||||
float barrier_mult = 1.0f;
|
||||
if (barrier_label != 0) {
|
||||
barrier_mult = 1.0f + barrier_scale * (float)barrier_label;
|
||||
}
|
||||
|
||||
float combined_reward = pnl_reward * barrier_mult
|
||||
+ diversity_scale * div_penalty
|
||||
+ curiosity_scale * curiosity_reward;
|
||||
|
||||
/* ---- Step 12: Episode done check ---- */
|
||||
step_in_episode++;
|
||||
int time_done = (step_in_episode >= episode_length) ? 1 : 0;
|
||||
int barrier_done = (barrier_label != 0) ? 1 : 0;
|
||||
int data_done = (next_bar >= total_bars) ? 1 : 0;
|
||||
int done = (time_done || barrier_done || data_done) ? 1 : 0;
|
||||
|
||||
/* ---- Step 13: Store per-timestep outputs ---- */
|
||||
for (int i = 0; i < STATE_DIM; i++)
|
||||
out_states[out_off * STATE_DIM + i] = state[i];
|
||||
out_actions[out_off] = action_idx;
|
||||
out_log_probs[out_off] = log_prob;
|
||||
out_dones[out_off] = done;
|
||||
|
||||
/* Store for GAE backward scan */
|
||||
gae_rewards[t] = combined_reward;
|
||||
gae_dones[t] = (float)done;
|
||||
|
||||
/* ---- Step 14: Episode reset on done ---- */
|
||||
if (done) {
|
||||
cash = initial_cap;
|
||||
position = 0.0f;
|
||||
entry_price = 0.0f;
|
||||
cum_costs = 0.0f;
|
||||
last_price = 0.0f;
|
||||
step_in_episode = 0;
|
||||
barrier_reset(barrier_st);
|
||||
/* Reset diversity window */
|
||||
for (int i = 0; i < DIVERSITY_WINDOW; i++)
|
||||
div_window[i] = 0;
|
||||
div_meta[0] = 0;
|
||||
div_meta[1] = 0;
|
||||
}
|
||||
} /* end Phase A timestep loop */
|
||||
|
||||
/* ================================================================ */
|
||||
/* Phase B: Backward GAE Scan */
|
||||
/* ================================================================ */
|
||||
|
||||
/* Step 1: Compute bootstrap value for GAE */
|
||||
int last_t = actual_L - 1;
|
||||
int last_out_off = tid * actual_L + last_t;
|
||||
if (out_dones[last_out_off] == 1) {
|
||||
/* Last step was terminal — bootstrap value is 0 */
|
||||
gae_values[actual_L] = 0.0f;
|
||||
} else {
|
||||
/* Last step was not terminal — run critic on next_state for bootstrap */
|
||||
int last_global_bar = ep_start + last_t;
|
||||
int next_bar_boot = last_global_bar + 1;
|
||||
|
||||
/* Build bootstrap next_state */
|
||||
float boot_state[STATE_DIM];
|
||||
if (next_bar_boot < total_bars) {
|
||||
int nmf_off = next_bar_boot * MARKET_DIM;
|
||||
for (int i = 0; i < MARKET_DIM; i++)
|
||||
boot_state[i] = market_features[nmf_off + i];
|
||||
} else {
|
||||
/* Reuse last state market features from output */
|
||||
for (int i = 0; i < MARKET_DIM; i++)
|
||||
boot_state[i] = out_states[last_out_off * STATE_DIM + i];
|
||||
}
|
||||
|
||||
/* Approximate portfolio features from current thread state */
|
||||
float boot_price_raw = 0.0f;
|
||||
if (next_bar_boot < total_bars) {
|
||||
int t_off_boot = next_bar_boot * 4;
|
||||
boot_price_raw = targets[t_off_boot + 2];
|
||||
if (boot_price_raw <= 0.0f) boot_price_raw = targets[t_off_boot + 0];
|
||||
}
|
||||
if (boot_price_raw <= 0.0f) boot_price_raw = 1.0f;
|
||||
|
||||
float boot_value = cash + position * boot_price_raw;
|
||||
float boot_norm = boot_value / initial_cap;
|
||||
float boot_max_pos = (boot_price_raw > 0.0f) ? initial_cap / boot_price_raw : 1.0f;
|
||||
|
||||
boot_state[MARKET_DIM + 0] = boot_norm;
|
||||
boot_state[MARKET_DIM + 1] = position / boot_max_pos;
|
||||
boot_state[MARKET_DIM + 2] = spread;
|
||||
|
||||
gae_values[actual_L] = ppo_critic_forward(
|
||||
boot_state,
|
||||
vw1, vb1, vw2, vb2, vw3, vb3,
|
||||
vw4, vb4, vw5, vb5, vw6, vb6,
|
||||
critic_a, critic_b
|
||||
);
|
||||
}
|
||||
|
||||
/* Step 2: Run GAE backward scan */
|
||||
compute_gae_backward(
|
||||
gae_rewards, gae_values, gae_dones,
|
||||
gae_advantages, gae_returns,
|
||||
actual_L, gamma, gae_lambda
|
||||
);
|
||||
|
||||
/* Step 3: Write advantages and returns to output buffers */
|
||||
for (int t = 0; t < actual_L; t++) {
|
||||
int out_off = tid * actual_L + t;
|
||||
out_advantages[out_off] = gae_advantages[t];
|
||||
out_returns[out_off] = gae_returns[t];
|
||||
}
|
||||
|
||||
/* ---- Write back per-thread state ---- */
|
||||
portfolio_states[ps_off + 0] = cash;
|
||||
portfolio_states[ps_off + 1] = position;
|
||||
portfolio_states[ps_off + 2] = entry_price;
|
||||
portfolio_states[ps_off + 3] = initial_cap;
|
||||
portfolio_states[ps_off + 4] = spread;
|
||||
portfolio_states[ps_off + 5] = last_price;
|
||||
portfolio_states[ps_off + 6] = reserve_pct;
|
||||
portfolio_states[ps_off + 7] = cum_costs;
|
||||
|
||||
for (int i = 0; i < BARRIER_STATE_SIZE; i++)
|
||||
barrier_states[bs_off + i] = barrier_st[i];
|
||||
|
||||
for (int i = 0; i < DIVERSITY_WINDOW; i++)
|
||||
diversity_windows[dw_off + i] = div_window[i];
|
||||
|
||||
diversity_metas[dm_off + 0] = div_meta[0];
|
||||
diversity_metas[dm_off + 1] = div_meta[1];
|
||||
|
||||
rng_states[tid] = rng;
|
||||
}
|
||||
Reference in New Issue
Block a user