feat: CVaR position scaling kernel wiring — env_step accepts cvar_scales
- experience_env_step kernel: new cvar_scales parameter (NULL = no scaling) - target_position *= cvar_scales[i] when buffer is non-NULL - GpuExperienceCollector: cvar_scales_ptr field + set_cvar_scales() setter - Default: NULL (0) = no scaling until IQN CVaR is wired from training loop The GpuIqnHead.compute_cvar_scales() produces the buffer, the collector passes it to the kernel. Full wiring through training_loop.rs is the remaining step — needs the collector to receive the device pointer from the fused training context after each IQN training step. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -386,7 +386,8 @@ extern "C" __global__ void experience_env_step(
|
||||
int b0_size,
|
||||
int b1_size,
|
||||
int b2_size,
|
||||
int current_t
|
||||
int current_t,
|
||||
const float* __restrict__ cvar_scales /* [N] or NULL — CVaR position scaling */
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= N) return;
|
||||
@@ -458,6 +459,11 @@ extern "C" __global__ void experience_env_step(
|
||||
float target_exposure = exposure_idx_to_fraction(exposure_idx);
|
||||
float target_position = target_exposure * max_position;
|
||||
|
||||
/* Apply CVaR risk scaling: reduce position when tail risk is high */
|
||||
if (cvar_scales != NULL) {
|
||||
target_position *= cvar_scales[i];
|
||||
}
|
||||
|
||||
/* ---- Position adjustment with volatility-scaled transaction cost ---- */
|
||||
/* Real spread widens in volatile markets. vol_scale (from CUSUM) is
|
||||
* computed later, but we can read CUSUM here for the cost calculation.
|
||||
|
||||
@@ -351,6 +351,11 @@ pub struct GpuExperienceCollector {
|
||||
/// Per-tensor sizes for computing weight pointers.
|
||||
param_sizes: [usize; 20],
|
||||
|
||||
// ── CVaR position scaling (IQN dual-head) ──────────────────────
|
||||
/// Device pointer to CVaR scales [N]. 0 = NULL = no scaling.
|
||||
/// Set by `set_cvar_scales()` before experience collection.
|
||||
cvar_scales_ptr: u64,
|
||||
|
||||
// ── Compiled experience kernels ─────────────────────────────────
|
||||
state_gather_kernel: CudaFunction,
|
||||
action_select_kernel: CudaFunction,
|
||||
@@ -770,6 +775,7 @@ impl GpuExperienceCollector {
|
||||
online_params_flat,
|
||||
total_params,
|
||||
param_sizes,
|
||||
cvar_scales_ptr: 0, // NULL = no CVaR scaling initially
|
||||
state_gather_kernel,
|
||||
action_select_kernel,
|
||||
env_step_kernel,
|
||||
@@ -816,6 +822,13 @@ impl GpuExperienceCollector {
|
||||
///
|
||||
/// Runs the timestep-level loop: for each timestep, gathers states,
|
||||
/// runs cuBLAS Q-forward, selects actions, and steps the environment.
|
||||
/// Set CVaR position scaling from IQN dual-head.
|
||||
/// The device pointer will be passed to the env_step kernel.
|
||||
/// Call with 0 to disable (NULL pointer = no scaling).
|
||||
pub fn set_cvar_scales(&mut self, device_ptr: u64) {
|
||||
self.cvar_scales_ptr = device_ptr;
|
||||
}
|
||||
|
||||
/// All outputs remain GPU-resident for zero-copy training.
|
||||
///
|
||||
/// The caller **must** call `stream().synchronize()` before passing
|
||||
@@ -1072,6 +1085,7 @@ impl GpuExperienceCollector {
|
||||
.arg(&b1)
|
||||
.arg(&b2)
|
||||
.arg(&t_i32)
|
||||
.arg(&self.cvar_scales_ptr) // CVaR position scaling (0 = NULL = no scaling)
|
||||
.launch(launch_cfg)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"experience_env_step t={t}: {e}"
|
||||
|
||||
Reference in New Issue
Block a user