feat(ensemble): add DQN and PPO inference adapters
Implement ModelInferenceAdapter for DQN (Q-value argmax -> direction + softmax confidence) and PPO (probability-weighted action values -> direction + max prob confidence) with zero-padding for mismatched feature dimensions and 6 passing unit tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
205
ml/src/ensemble/adapters/dqn.rs
Normal file
205
ml/src/ensemble/adapters/dqn.rs
Normal file
@@ -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<DQN>,
|
||||
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<Self> {
|
||||
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<Self> {
|
||||
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<EnsemblePrediction> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Convert f64 feature values to f32 for candle tensor
|
||||
let f32_values: Vec<f32> = 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<f64>
|
||||
let q_squeezed = q_output
|
||||
.squeeze(0)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze Q-values: {e}")))?;
|
||||
let q_vec_f32: Vec<f32> = q_squeezed
|
||||
.to_vec1()
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to extract Q-values: {e}")))?;
|
||||
let q_vec: Vec<f64> = 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
7
ml/src/ensemble/adapters/mod.rs
Normal file
7
ml/src/ensemble/adapters/mod.rs
Normal file
@@ -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;
|
||||
206
ml/src/ensemble/adapters/ppo.rs
Normal file
206
ml/src/ensemble/adapters/ppo.rs
Normal file
@@ -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<PPO>,
|
||||
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<Self> {
|
||||
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<f32> {
|
||||
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<EnsemblePrediction> {
|
||||
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<f32> = 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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::{
|
||||
|
||||
Reference in New Issue
Block a user