diff --git a/ml/src/ensemble/adapters/dqn.rs b/ml/src/ensemble/adapters/dqn.rs new file mode 100644 index 000000000..cac1d993e --- /dev/null +++ b/ml/src/ensemble/adapters/dqn.rs @@ -0,0 +1,205 @@ +//! DQN inference adapter for ensemble prediction +//! +//! Wraps a loaded DQN model and normalizes its Q-value output +//! into a directional signal + confidence for ensemble aggregation. + +use std::sync::Mutex; + +use candle_core::{Device, Tensor}; + +use crate::dqn::dqn::{DQNConfig, DQN}; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// Inference adapter that wraps a DQN model for ensemble prediction. +/// +/// Converts raw Q-values into a normalized directional signal [-1, 1] +/// and a confidence score [0, 1] based on softmax probability of the +/// best action. +#[allow(missing_debug_implementations)] +pub struct DqnInferenceAdapter { + model: Mutex, + device: Device, +} + +// SAFETY: DQN internally uses candle tensors which are Send+Sync. +// The Mutex provides exclusive access for inference calls. +unsafe impl Send for DqnInferenceAdapter {} +unsafe impl Sync for DqnInferenceAdapter {} + +impl DqnInferenceAdapter { + /// Create a new DQN inference adapter from configuration. + /// + /// Initializes a fresh DQN model with random weights on the best + /// available device (CUDA GPU if available, otherwise CPU). + pub fn new(config: DQNConfig) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let model = DQN::new(config)?; + Ok(Self { + model: Mutex::new(model), + device, + }) + } + + /// Create a DQN inference adapter and load weights from a safetensors checkpoint. + pub fn from_checkpoint(config: DQNConfig, path: &str) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let mut model = DQN::new(config)?; + model.load_from_safetensors(path)?; + Ok(Self { + model: Mutex::new(model), + device, + }) + } +} + +impl ModelInferenceAdapter for DqnInferenceAdapter { + fn model_name(&self) -> &str { + "DQN" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + // Convert f64 feature values to f32 for candle tensor + let f32_values: Vec = features.values.iter().map(|&v| v as f32).collect(); + let len = f32_values.len(); + + // Create input tensor [1, feature_dim] + let input = Tensor::from_vec(f32_values, (1, len), &self.device) + .map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {e}")))?; + + // Run forward pass through the Q-network + let model = self + .model + .lock() + .map_err(|e| MLError::LockError(format!("DQN model lock poisoned: {e}")))?; + let q_output = model.forward(&input)?; + + // Squeeze batch dimension and extract Q-values as Vec + let q_squeezed = q_output + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("Failed to squeeze Q-values: {e}")))?; + let q_vec_f32: Vec = q_squeezed + .to_vec1() + .map_err(|e| MLError::ModelError(format!("Failed to extract Q-values: {e}")))?; + let q_vec: Vec = q_vec_f32.iter().map(|&v| v as f64).collect(); + + let num_actions = q_vec.len(); + if num_actions == 0 { + return Err(MLError::InferenceError( + "DQN produced zero-length Q-value vector".to_string(), + )); + } + + // Find best action index (argmax of Q-values) + let best_idx = q_vec + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(idx, _)| idx) + .ok_or_else(|| MLError::InferenceError("Failed to find max Q-value".to_string()))?; + + // Compute directional signal: map action index to [-1, 1] + // center = (num_actions - 1) / 2.0 + // direction = (best_idx - center) / center, clamped to [-1, 1] + let center = (num_actions as f64 - 1.0) / 2.0; + let direction = if center > 0.0 { + ((best_idx as f64 - center) / center).clamp(-1.0, 1.0) + } else { + 0.0 + }; + + // Compute confidence via softmax probability of the best action + let max_q = q_vec + .iter() + .copied() + .fold(f64::NEG_INFINITY, f64::max); + let exp_sum: f64 = q_vec.iter().map(|&q| (q - max_q).exp()).sum(); + let best_q = q_vec + .get(best_idx) + .copied() + .ok_or_else(|| MLError::InferenceError("Best index out of bounds".to_string()))?; + let confidence = (best_q - max_q).exp() / exp_sum; + + let latency_us = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "DQN".to_string(), + direction, + confidence, + metadata: PredictionMeta { + latency_us, + quantiles: None, + attention_weights: None, + q_values: Some(q_vec), + }, + }) + } + + fn is_ready(&self) -> bool { + self.model.lock().is_ok() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_config() -> DQNConfig { + DQNConfig { + state_dim: 51, + num_actions: 45, + hidden_dims: vec![64, 64], + ..Default::default() + } + } + + #[test] + fn test_dqn_adapter_creation() { + let adapter = DqnInferenceAdapter::new(test_config()).unwrap(); + assert_eq!(adapter.model_name(), "DQN"); + assert!(adapter.is_ready()); + } + + #[test] + fn test_dqn_adapter_predict_direction_range() { + let adapter = DqnInferenceAdapter::new(test_config()).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred = adapter.predict(&fv).unwrap(); + assert!( + pred.direction >= -1.0 && pred.direction <= 1.0, + "direction {} out of [-1,1]", + pred.direction + ); + assert!( + pred.confidence >= 0.0 && pred.confidence <= 1.0, + "confidence {} out of [0,1]", + pred.confidence + ); + assert!( + pred.metadata.q_values.is_some(), + "q_values metadata should be Some" + ); + } + + #[test] + fn test_dqn_adapter_deterministic() { + let adapter = DqnInferenceAdapter::new(test_config()).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred1 = adapter.predict(&fv).unwrap(); + let pred2 = adapter.predict(&fv).unwrap(); + assert_eq!( + pred1.direction, pred2.direction, + "Deterministic predictions should have same direction" + ); + } +} diff --git a/ml/src/ensemble/adapters/mod.rs b/ml/src/ensemble/adapters/mod.rs new file mode 100644 index 000000000..e653181ab --- /dev/null +++ b/ml/src/ensemble/adapters/mod.rs @@ -0,0 +1,7 @@ +//! Per-model inference adapters for ensemble prediction + +pub mod dqn; +pub mod ppo; + +pub use dqn::DqnInferenceAdapter; +pub use ppo::PpoInferenceAdapter; diff --git a/ml/src/ensemble/adapters/ppo.rs b/ml/src/ensemble/adapters/ppo.rs new file mode 100644 index 000000000..b6cbc8887 --- /dev/null +++ b/ml/src/ensemble/adapters/ppo.rs @@ -0,0 +1,206 @@ +//! PPO inference adapter for ensemble prediction +//! +//! Wraps a loaded PPO model (MLP actor) and normalizes its action +//! probability output into a directional signal + confidence for +//! ensemble aggregation. + +use std::sync::Mutex; + +use candle_core::{Device, Tensor}; + +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::ppo::ppo::{PPOConfig, PPO}; +use crate::{MLError, MLResult}; + +/// Inference adapter that wraps a PPO model for ensemble prediction. +/// +/// Converts the actor's softmax action probabilities into a normalized +/// directional signal [-1, 1] using probability-weighted action values, +/// and uses the max probability as confidence [0, 1]. +#[allow(missing_debug_implementations)] +pub struct PpoInferenceAdapter { + model: Mutex, + state_dim: usize, + device: Device, +} + +// SAFETY: PPO internally uses candle tensors which are Send+Sync. +// The Mutex provides exclusive access for inference calls. +unsafe impl Send for PpoInferenceAdapter {} +unsafe impl Sync for PpoInferenceAdapter {} + +impl PpoInferenceAdapter { + /// Create a new PPO inference adapter from configuration. + /// + /// Initializes a fresh PPO model with random weights on the best + /// available device (CUDA GPU if available, otherwise CPU). + pub fn new(config: PPOConfig) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let state_dim = config.state_dim; + let model = PPO::new(config)?; + Ok(Self { + model: Mutex::new(model), + state_dim, + device, + }) + } + + /// Pad or truncate feature values to match the expected state_dim. + /// + /// Takes up to `state_dim` values from the input, converting f64 to f32. + /// If the input is shorter than `state_dim`, the remaining elements are + /// zero-padded. + fn pad_features(&self, values: &[f64]) -> Vec { + let mut padded = vec![0.0f32; self.state_dim]; + let copy_len = values.len().min(self.state_dim); + for i in 0..copy_len { + if let Some(v) = values.get(i) { + if let Some(slot) = padded.get_mut(i) { + *slot = *v as f32; + } + } + } + padded + } +} + +impl ModelInferenceAdapter for PpoInferenceAdapter { + fn model_name(&self) -> &str { + "PPO" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + // Pad features to state_dim, converting f64 -> f32 + let padded = self.pad_features(&features.values); + + // Create input tensor [1, state_dim] + let input = Tensor::from_vec(padded, (1, self.state_dim), &self.device) + .map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {e}")))?; + + // Get action probabilities from the actor network (softmax output) + let model = self + .model + .lock() + .map_err(|e| MLError::LockError(format!("PPO model lock poisoned: {e}")))?; + let probs_tensor = model.actor.action_probabilities(&input)?; + + // Squeeze batch dimension and extract probabilities + let probs_squeezed = probs_tensor + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("Failed to squeeze probabilities: {e}")))?; + let probs: Vec = probs_squeezed + .to_vec1() + .map_err(|e| MLError::ModelError(format!("Failed to extract probabilities: {e}")))?; + + let num_actions = probs.len(); + if num_actions == 0 { + return Err(MLError::InferenceError( + "PPO produced zero-length probability vector".to_string(), + )); + } + + // Compute directional signal: weighted sum of probs * centered action values + // Action values are centered: action_value[i] = (i - center) / center + let center = (num_actions as f64 - 1.0) / 2.0; + let direction = if center > 0.0 { + let weighted_sum: f64 = probs + .iter() + .enumerate() + .map(|(i, &p)| { + let action_val = (i as f64 - center) / center; + p as f64 * action_val + }) + .sum(); + weighted_sum.clamp(-1.0, 1.0) + } else { + 0.0 + }; + + // Confidence: max probability across all actions + let confidence = probs + .iter() + .copied() + .fold(0.0f32, f32::max) as f64; + + let latency_us = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "PPO".to_string(), + direction, + confidence: confidence.clamp(0.0, 1.0), + metadata: PredictionMeta { + latency_us, + quantiles: None, + attention_weights: None, + q_values: None, + }, + }) + } + + fn is_ready(&self) -> bool { + self.model.lock().is_ok() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_config() -> PPOConfig { + PPOConfig { + state_dim: 64, + num_actions: 45, + policy_hidden_dims: vec![64, 64], + value_hidden_dims: vec![64, 64], + ..Default::default() + } + } + + #[test] + fn test_ppo_adapter_creation() { + let adapter = PpoInferenceAdapter::new(test_config()).unwrap(); + assert_eq!(adapter.model_name(), "PPO"); + assert!(adapter.is_ready()); + assert_eq!(adapter.state_dim, 64); + } + + #[test] + fn test_ppo_adapter_pads_input_to_64() { + let adapter = PpoInferenceAdapter::new(test_config()).unwrap(); + // 51-dim input (DQN canonical size) should be zero-padded to 64 + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred = adapter.predict(&fv).unwrap(); + assert!( + pred.direction >= -1.0 && pred.direction <= 1.0, + "direction {} out of [-1,1]", + pred.direction + ); + assert!( + pred.confidence >= 0.0 && pred.confidence <= 1.0, + "confidence {} out of [0,1]", + pred.confidence + ); + } + + #[test] + fn test_ppo_adapter_deterministic() { + let adapter = PpoInferenceAdapter::new(test_config()).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred1 = adapter.predict(&fv).unwrap(); + let pred2 = adapter.predict(&fv).unwrap(); + assert_eq!( + pred1.direction, pred2.direction, + "Deterministic predictions should have same direction" + ); + } +} diff --git a/ml/src/ensemble/mod.rs b/ml/src/ensemble/mod.rs index 1861a8011..18a087888 100644 --- a/ml/src/ensemble/mod.rs +++ b/ml/src/ensemble/mod.rs @@ -18,6 +18,7 @@ pub mod training_integration; // Training integration for ML service pub mod voting; pub mod weights; pub mod inference_adapter; +pub mod adapters; // Re-export key types that are used across ensemble modules pub use ab_testing::{