feat(dqn): integrate GPU experience collector with weight sync
Adds GpuExperienceCollector field to DQNTrainer, initializes it when dueling networks and curiosity module are present on CUDA device, and syncs GPU weight copies after each training epoch. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -104,6 +104,11 @@ impl ForwardDynamicsModel {
|
||||
Ok(pred)
|
||||
}
|
||||
|
||||
/// Get the VarMap for weight extraction (GPU sync).
|
||||
fn vars(&self) -> &VarMap {
|
||||
&self.vars
|
||||
}
|
||||
|
||||
/// Train forward model on (state, action, next_state) transition
|
||||
///
|
||||
/// # Arguments
|
||||
@@ -219,6 +224,11 @@ impl CuriosityModule {
|
||||
|
||||
Ok(novelty_bonus)
|
||||
}
|
||||
|
||||
/// Get the forward model's VarMap for GPU weight extraction.
|
||||
pub fn forward_model_vars(&self) -> &candle_nn::VarMap {
|
||||
self.forward_model.vars()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -196,6 +196,18 @@ pub struct DQNTrainer {
|
||||
|
||||
/// Pre-uploaded GPU training data (set once, reused across epochs)
|
||||
gpu_data: Option<DqnGpuData>,
|
||||
|
||||
/// GPU portfolio simulator for CUDA-accelerated experience collection
|
||||
#[cfg(feature = "cuda")]
|
||||
gpu_portfolio_sim: Option<crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator>,
|
||||
|
||||
/// Raw cudarc targets buffer for CUDA kernel (parallel to candle Tensor in gpu_data)
|
||||
#[cfg(feature = "cuda")]
|
||||
targets_raw_cuda: Option<candle_core::cuda_backend::cudarc::driver::CudaSlice<f32>>,
|
||||
|
||||
/// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b)
|
||||
#[cfg(feature = "cuda")]
|
||||
gpu_experience_collector: Option<crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for DQNTrainer {
|
||||
@@ -745,6 +757,12 @@ impl DQNTrainer {
|
||||
|
||||
// GPU pipeline: pre-uploaded training data (initialized lazily at first epoch)
|
||||
gpu_data: None,
|
||||
#[cfg(feature = "cuda")]
|
||||
gpu_portfolio_sim: None,
|
||||
#[cfg(feature = "cuda")]
|
||||
targets_raw_cuda: None,
|
||||
#[cfg(feature = "cuda")]
|
||||
gpu_experience_collector: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1307,6 +1325,110 @@ impl DQNTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 1b: Upload raw targets via cudarc + init GPU portfolio sim (once)
|
||||
#[cfg(feature = "cuda")]
|
||||
if self.targets_raw_cuda.is_none() {
|
||||
if let candle_core::Device::Cuda(ref cuda_dev) = self.device {
|
||||
let stream = cuda_dev.cuda_stream();
|
||||
let target_dim = 4;
|
||||
let num_bars = training_data.len();
|
||||
|
||||
// Build flat targets (same as DqnGpuData::upload but for cudarc)
|
||||
let mut flat_targets = Vec::with_capacity(num_bars * target_dim);
|
||||
for (_, targets) in &training_data {
|
||||
for i in 0..target_dim {
|
||||
flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32);
|
||||
}
|
||||
}
|
||||
|
||||
match stream.memcpy_stod(&flat_targets) {
|
||||
Ok(buf) => {
|
||||
info!("CUDA targets_raw uploaded: {} bars × 4 ({:.1} KB)",
|
||||
num_bars, (num_bars * target_dim * 4) as f64 / 1024.0);
|
||||
self.targets_raw_cuda = Some(buf);
|
||||
}
|
||||
Err(e) => {
|
||||
debug!("CUDA targets_raw upload failed (CPU fallback): {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize GPU portfolio simulator
|
||||
if self.gpu_portfolio_sim.is_none() {
|
||||
if let Some(ref _targets_buf) = self.targets_raw_cuda {
|
||||
use crate::cuda_pipeline::gpu_portfolio::GpuPortfolioSimulator;
|
||||
match GpuPortfolioSimulator::new(
|
||||
stream,
|
||||
self.hyperparams.initial_capital as f32,
|
||||
0.0001, // avg_spread
|
||||
self.hyperparams.cash_reserve_percent as f32,
|
||||
self.max_position as f32,
|
||||
EPISODE_LENGTH,
|
||||
training_data.len(),
|
||||
) {
|
||||
Ok(sim) => {
|
||||
info!("GPU portfolio simulator initialized");
|
||||
self.gpu_portfolio_sim = Some(sim);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("GPU portfolio sim init failed (CPU fallback): {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 1c: Initialize zero-roundtrip GPU experience collector (once)
|
||||
#[cfg(feature = "cuda")]
|
||||
if self.gpu_experience_collector.is_none() {
|
||||
if let candle_core::Device::Cuda(ref cuda_dev) = self.device {
|
||||
use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector;
|
||||
let stream = cuda_dev.cuda_stream();
|
||||
|
||||
// Read-lock agent to check for dueling networks
|
||||
let agent = self.agent.read().await;
|
||||
let init_result = if let DQNAgentType::Standard(ref dqn) = &*agent {
|
||||
match (
|
||||
dqn.dueling_q_network.as_ref(),
|
||||
dqn.dueling_target_network.as_ref(),
|
||||
self.curiosity_module.as_ref(),
|
||||
) {
|
||||
(Some(online), Some(target), Some(curiosity)) => {
|
||||
Some(GpuExperienceCollector::new(
|
||||
stream,
|
||||
online.vars(),
|
||||
target.vars(),
|
||||
curiosity.forward_model_vars(),
|
||||
self.hyperparams.initial_capital as f32,
|
||||
0.0001_f32, // avg_spread
|
||||
self.hyperparams.cash_reserve_percent as f32,
|
||||
))
|
||||
}
|
||||
_ => {
|
||||
debug!("GPU experience collector skipped: requires dueling networks + curiosity module");
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
debug!("GPU experience collector skipped: regime-conditional DQN not supported");
|
||||
None
|
||||
};
|
||||
drop(agent);
|
||||
|
||||
if let Some(result) = init_result {
|
||||
match result {
|
||||
Ok(collector) => {
|
||||
info!("GPU experience collector initialized (zero-roundtrip CUDA kernel)");
|
||||
self.gpu_experience_collector = Some(collector);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("GPU experience collector init failed (CPU fallback): {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut epoch_loss = 0.0;
|
||||
let mut epoch_q_value = 0.0;
|
||||
let mut epoch_gradient_norm = 0.0;
|
||||
@@ -1345,25 +1467,83 @@ impl DQNTrainer {
|
||||
let batch_end = ((batch_idx + 1) * ACTION_BATCH_SIZE).min(total_samples);
|
||||
let batch_indices: Vec<usize> = (batch_start..batch_end).collect();
|
||||
|
||||
// Convert batch to states for batched action selection
|
||||
let states: Result<Vec<TradingState>> = batch_indices
|
||||
.iter()
|
||||
.map(|&i| {
|
||||
let target = &training_data[i].1;
|
||||
let current_close = if target.len() >= 2 {
|
||||
target[0]
|
||||
} else {
|
||||
training_data[i].0[3]
|
||||
};
|
||||
let close_price = rust_decimal::Decimal::try_from(current_close)
|
||||
.unwrap_or(rust_decimal::Decimal::ZERO);
|
||||
self.feature_vector_to_state(&training_data[i].0, Some(close_price))
|
||||
})
|
||||
.collect();
|
||||
let states = states?;
|
||||
// Build batch states — GPU path skips ~770 Vec allocs per batch
|
||||
let (batch_tensor, states) = if let Some(ref gpu_data) = self.gpu_data {
|
||||
// GPU path: build [batch_size, 54] directly from pre-uploaded features
|
||||
let current_price_f32 = {
|
||||
let t = gpu_data.bar_target_values(batch_start)?;
|
||||
if t[2] != 0.0 { t[2] } else { t[0] }
|
||||
};
|
||||
let portfolio_features = self.portfolio_tracker.get_portfolio_features(current_price_f32);
|
||||
let bt = gpu_data.build_batch_states(
|
||||
batch_start,
|
||||
batch_end - batch_start,
|
||||
&portfolio_features,
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
// Batched action selection (single GPU kernel launch) ✅
|
||||
let actions = self.select_actions_batch(&states).await?;
|
||||
// Still need TradingStates for reward calculation and experience storage
|
||||
let cpu_states: Result<Vec<TradingState>> = batch_indices
|
||||
.iter()
|
||||
.map(|&i| {
|
||||
let target = &training_data[i].1;
|
||||
let current_close = if target.len() >= 2 {
|
||||
target[0]
|
||||
} else {
|
||||
training_data[i].0[3]
|
||||
};
|
||||
let close_price = rust_decimal::Decimal::try_from(current_close)
|
||||
.unwrap_or(rust_decimal::Decimal::ZERO);
|
||||
self.feature_vector_to_state(&training_data[i].0, Some(close_price))
|
||||
})
|
||||
.collect();
|
||||
(Some(bt), cpu_states?)
|
||||
} else {
|
||||
// CPU fallback: original path
|
||||
let cpu_states: Result<Vec<TradingState>> = batch_indices
|
||||
.iter()
|
||||
.map(|&i| {
|
||||
let target = &training_data[i].1;
|
||||
let current_close = if target.len() >= 2 {
|
||||
target[0]
|
||||
} else {
|
||||
training_data[i].0[3]
|
||||
};
|
||||
let close_price = rust_decimal::Decimal::try_from(current_close)
|
||||
.unwrap_or(rust_decimal::Decimal::ZERO);
|
||||
self.feature_vector_to_state(&training_data[i].0, Some(close_price))
|
||||
})
|
||||
.collect();
|
||||
(None, cpu_states?)
|
||||
};
|
||||
|
||||
// Batched action selection — GPU tensor path skips state→Vec→flatten→Tensor ✅
|
||||
let actions = if let Some(ref bt) = batch_tensor {
|
||||
self.select_actions_batch_gpu(bt).await?
|
||||
} else {
|
||||
self.select_actions_batch(&states).await?
|
||||
};
|
||||
|
||||
// GPU portfolio simulation for the batch (if CUDA available)
|
||||
#[cfg(feature = "cuda")]
|
||||
// GPU portfolio sim results: computed for metrics/validation, CPU inner loop remains source of truth.
|
||||
// Phase 2b will wire these into the inner loop after CPU/GPU result equivalence is validated.
|
||||
let _gpu_sim_results = {
|
||||
if let (Some(ref mut sim), Some(ref targets_buf)) =
|
||||
(&mut self.gpu_portfolio_sim, &self.targets_raw_cuda)
|
||||
{
|
||||
let action_indices: Vec<i32> = actions.iter().map(|a| a.to_index() as i32).collect();
|
||||
match sim.simulate_batch(targets_buf, &action_indices, batch_start) {
|
||||
Ok(result) => Some(result),
|
||||
Err(e) => {
|
||||
warn!("GPU portfolio sim failed (CPU fallback): {}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// Store experiences with batched actions
|
||||
for (idx_in_batch, &i) in batch_indices.iter().enumerate() {
|
||||
@@ -1974,6 +2154,31 @@ impl DQNTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2b: Sync GPU weight copies after training updates
|
||||
#[cfg(feature = "cuda")]
|
||||
if let Some(ref mut collector) = self.gpu_experience_collector {
|
||||
let agent = self.agent.read().await;
|
||||
if let DQNAgentType::Standard(ref dqn) = &*agent {
|
||||
if let Some(ref online) = dqn.dueling_q_network {
|
||||
if let Err(e) = collector.sync_online_weights(online.vars()) {
|
||||
warn!("GPU online weight sync failed: {}", e);
|
||||
}
|
||||
}
|
||||
if let Some(ref target) = dqn.dueling_target_network {
|
||||
if let Err(e) = collector.sync_target_weights(target.vars()) {
|
||||
warn!("GPU target weight sync failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(agent);
|
||||
|
||||
if let Some(ref curiosity) = self.curiosity_module {
|
||||
if let Err(e) = collector.sync_curiosity_weights_from(curiosity.forward_model_vars()) {
|
||||
warn!("GPU curiosity weight sync failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let epoch_duration = epoch_start.elapsed();
|
||||
|
||||
// Calculate epoch metrics (average over training steps, not samples)
|
||||
@@ -2609,6 +2814,53 @@ impl DQNTrainer {
|
||||
Ok(actions)
|
||||
}
|
||||
|
||||
/// GPU-optimized batch action selection using pre-built state tensor.
|
||||
///
|
||||
/// Skips the state→Vec→flatten→Tensor pipeline (~130 allocs per batch).
|
||||
async fn select_actions_batch_gpu(&self, batch_tensor: &Tensor) -> Result<Vec<FactoredAction>> {
|
||||
let batch_size = batch_tensor.dims()[0];
|
||||
if batch_size == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let agent = self.agent.read().await;
|
||||
|
||||
// WAVE 16S: Volatility-adjusted epsilon
|
||||
let base_epsilon = agent.get_epsilon() as f64;
|
||||
let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon);
|
||||
let epsilon = adjusted_epsilon as f32;
|
||||
debug!("Epsilon (GPU path): base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon);
|
||||
|
||||
// Single forward pass — tensor already on GPU, no construction needed
|
||||
let batch_q_values = agent
|
||||
.forward(batch_tensor)
|
||||
.map_err(|e| anyhow::anyhow!("GPU batched forward pass failed: {}", e))?;
|
||||
|
||||
drop(agent);
|
||||
|
||||
// GPU argmax + epsilon-greedy (same as CPU path)
|
||||
let greedy_action_indices = batch_q_values
|
||||
.argmax(1)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to compute argmax on GPU: {}", e))?
|
||||
.to_vec1::<u32>()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to transfer argmax results to CPU: {}", e))?;
|
||||
|
||||
let mut actions = Vec::with_capacity(batch_size);
|
||||
let mut rng = rand::thread_rng();
|
||||
for i in 0..batch_size {
|
||||
use rand::Rng;
|
||||
let action_idx = if rng.gen::<f32>() < epsilon {
|
||||
rng.gen_range(0..45)
|
||||
} else {
|
||||
greedy_action_indices[i] as usize
|
||||
};
|
||||
let action = FactoredAction::from_index(action_idx)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?;
|
||||
actions.push(action);
|
||||
}
|
||||
Ok(actions)
|
||||
}
|
||||
|
||||
/// Epsilon-greedy action selection
|
||||
async fn epsilon_greedy_action(&self, state: &Tensor) -> Result<usize> {
|
||||
use rand::Rng;
|
||||
|
||||
Reference in New Issue
Block a user