diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 9f72d0555..095e46c6b 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -4116,17 +4116,26 @@ impl DQNTrainer { let action_indices_tensor = selector .select_actions(&batch_q_values, epsilon, batch_size, 5) .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))?; - // Single GPU->CPU readback for route_action() dispatch - let action_indices = action_indices_tensor + + // GPU-resident action routing: exposure (0-4) → factored (0-44) on-device. + // Single batch readback of factored indices — no per-item CPU routing. + let factored_tensor = selector.route_exposure_to_factored( + &action_indices_tensor, + batch_size, + self.hyperparams.avg_spread as f32, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))?; + + let factored_indices = factored_tensor .to_vec1::() - .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + .map_err(|e| anyhow::anyhow!("Factored indices readback: {e}"))?; let mut actions = Vec::with_capacity(batch_size); - for idx in &action_indices { - let action_idx = *idx as usize; - let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) - .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); + for idx in factored_indices { + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; actions.push(action); } return Ok(actions); @@ -4307,30 +4316,34 @@ impl DQNTrainer { .map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {}", e))? }; - let action_indices = action_indices_tensor + // GPU-resident action routing: exposure (0-4) → factored (0-44) on-device. + // Branching already produces factored indices, so pass through. + // Single batch readback of factored indices — no per-item CPU routing. + let factored_tensor = if self.hyperparams.use_branching { + // Branching: already factored (0-44), pass through + action_indices_tensor + } else { + // Route exposure (0-4) → factored (0-44) on GPU + selector.route_exposure_to_factored( + &action_indices_tensor, + batch_size, + self.hyperparams.avg_spread as f32, + self.hyperparams.avg_spread as f32, + self.vol_ema as f32, + self.median_vol as f32, + ).map_err(|e| anyhow::anyhow!("GPU route exposure→factored: {e}"))? + }; + + // Single batch readback (experience collection path, not training hot path) + let factored_indices = factored_tensor .to_vec1::() - .map_err(|e| anyhow::anyhow!("Failed to readback action indices: {}", e))?; + .map_err(|e| anyhow::anyhow!("Factored indices readback: {e}"))?; let mut actions = Vec::with_capacity(batch_size); - if self.hyperparams.use_branching { - // Branching: factored indices 0-44 - for &idx in &action_indices { - let action = FactoredAction::from_index(idx as usize) - .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; - actions.push(action); - } - } else { - // Standard: exposure indices 0-4 - for &idx in &action_indices { - let action_idx = idx as usize; - let exposure = crate::dqn::action_space::ExposureLevel::from_index(action_idx) - .map_err(|e| anyhow::anyhow!("Invalid exposure index {}: {}", action_idx, e))?; - // When routed kernel was used, exposure is already post-fill - // (unfilled → Flat override done on GPU). Route to get the - // FactoredAction struct but the exposure level is final. - let action = self.route_action(exposure, self.hyperparams.avg_spread as f32); - actions.push(action); - } + for idx in factored_indices { + let action = FactoredAction::from_index(idx as usize) + .map_err(|e| anyhow::anyhow!("Invalid factored index {}: {}", idx, e))?; + actions.push(action); } return Ok((actions, use_routed)); }