516 lines
17 KiB
Rust
516 lines
17 KiB
Rust
//! Ensemble-level hyperopt adapter
|
|
//!
|
|
//! Adds N ensemble weight parameters (one per model) to a combined parameter space,
|
|
//! enabling joint optimization of model weights alongside per-model hyperparameters.
|
|
//!
|
|
//! ## Design
|
|
//!
|
|
//! The [`ParameterSpace`] trait uses static methods (`continuous_bounds()`, `param_names()`),
|
|
//! which means the dimensionality must be known at compile time. To support variable-size
|
|
//! ensemble configurations at runtime, this module provides:
|
|
//!
|
|
//! - [`EnsembleSpaceConfig`]: Runtime configuration specifying model names and per-model
|
|
//! parameter dimensions. Stored in a thread-local so that static trait methods can
|
|
//! access it.
|
|
//! - [`EnsembleParameterSpace`]: The parameter space struct implementing [`ParameterSpace`].
|
|
//! Holds a flat continuous vector (per-model params concatenated + N weight values).
|
|
//!
|
|
//! ## Weight Normalization
|
|
//!
|
|
//! Raw weight values in `[0, 1]` are normalized via softmax-style division so they
|
|
//! sum to 1.0. If all raw weights are zero, equal weights are assigned.
|
|
//!
|
|
//! ## Usage
|
|
//!
|
|
//! ```rust,no_run
|
|
//! use ml::hyperopt::adapters::ensemble::{EnsembleSpaceConfig, EnsembleParameterSpace};
|
|
//! use ml::hyperopt::ParameterSpace;
|
|
//!
|
|
//! // Configure for a 3-model ensemble (DQN=11D, PPO=5D, TFT=6D)
|
|
//! let config = EnsembleSpaceConfig::new(
|
|
//! vec!["dqn".into(), "ppo".into(), "tft".into()],
|
|
//! vec![11, 5, 6],
|
|
//! // Per-model bounds: 11 DQN bounds + 5 PPO bounds + 6 TFT bounds
|
|
//! vec![
|
|
//! // ... 22 bounds total from individual model ParameterSpaces
|
|
//! ],
|
|
//! // Per-model param names
|
|
//! vec![
|
|
//! // ... 22 names total
|
|
//! ],
|
|
//! );
|
|
//!
|
|
//! // Install config so ParameterSpace static methods can read it
|
|
//! config.install();
|
|
//!
|
|
//! // Now EnsembleParameterSpace implements ParameterSpace
|
|
//! let bounds = EnsembleParameterSpace::continuous_bounds();
|
|
//! // bounds.len() == 22 (model params) + 3 (weights) = 25
|
|
//! ```
|
|
|
|
use std::cell::RefCell;
|
|
|
|
use crate::hyperopt::traits::ParameterSpace;
|
|
use crate::MLError;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Runtime configuration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Runtime configuration for the ensemble parameter space.
|
|
///
|
|
/// Must be [`install()`](EnsembleSpaceConfig::install)ed before calling
|
|
/// `EnsembleParameterSpace` trait methods.
|
|
#[derive(Debug, Clone)]
|
|
pub struct EnsembleSpaceConfig {
|
|
/// Human-readable model names (e.g., "dqn", "ppo", "tft").
|
|
pub model_names: Vec<String>,
|
|
/// Number of continuous parameters per model.
|
|
pub model_param_dims: Vec<usize>,
|
|
/// Concatenated per-model bounds `[(min, max), ...]`.
|
|
/// Length must equal `model_param_dims.iter().sum()`.
|
|
pub model_bounds: Vec<(f64, f64)>,
|
|
/// Concatenated per-model parameter names.
|
|
/// Length must equal `model_param_dims.iter().sum()`.
|
|
pub model_param_names: Vec<String>,
|
|
}
|
|
|
|
thread_local! {
|
|
static ENSEMBLE_CONFIG: RefCell<Option<EnsembleSpaceConfig>> = const { RefCell::new(None) };
|
|
}
|
|
|
|
impl EnsembleSpaceConfig {
|
|
/// Create a new ensemble space configuration.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `model_names` - Names of models in the ensemble.
|
|
/// * `model_param_dims` - Number of hyperparameters per model.
|
|
/// * `model_bounds` - Concatenated bounds for all model parameters.
|
|
/// * `model_param_names` - Concatenated parameter names for all models.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics if lengths are inconsistent.
|
|
pub fn new(
|
|
model_names: Vec<String>,
|
|
model_param_dims: Vec<usize>,
|
|
model_bounds: Vec<(f64, f64)>,
|
|
model_param_names: Vec<String>,
|
|
) -> Self {
|
|
let total_model_params: usize = model_param_dims.iter().sum();
|
|
assert_eq!(
|
|
model_names.len(),
|
|
model_param_dims.len(),
|
|
"model_names and model_param_dims must have the same length"
|
|
);
|
|
assert_eq!(
|
|
model_bounds.len(),
|
|
total_model_params,
|
|
"model_bounds length must equal sum of model_param_dims"
|
|
);
|
|
assert_eq!(
|
|
model_param_names.len(),
|
|
total_model_params,
|
|
"model_param_names length must equal sum of model_param_dims"
|
|
);
|
|
Self {
|
|
model_names,
|
|
model_param_dims,
|
|
model_bounds,
|
|
model_param_names,
|
|
}
|
|
}
|
|
|
|
/// Install this configuration in the thread-local slot.
|
|
///
|
|
/// Must be called before using `EnsembleParameterSpace` trait methods.
|
|
pub fn install(&self) {
|
|
ENSEMBLE_CONFIG.with(|cell| {
|
|
*cell.borrow_mut() = Some(self.clone());
|
|
});
|
|
}
|
|
|
|
/// Remove the installed configuration.
|
|
pub fn uninstall() {
|
|
ENSEMBLE_CONFIG.with(|cell| {
|
|
*cell.borrow_mut() = None;
|
|
});
|
|
}
|
|
|
|
/// Total dimension = sum(model_param_dims) + num_models (weight params).
|
|
pub fn total_dim(&self) -> usize {
|
|
let model_params: usize = self.model_param_dims.iter().sum();
|
|
model_params + self.model_names.len()
|
|
}
|
|
|
|
/// Number of models in the ensemble.
|
|
pub fn num_models(&self) -> usize {
|
|
self.model_names.len()
|
|
}
|
|
}
|
|
|
|
/// Read the installed config, returning an error if none is installed.
|
|
fn with_config<T>(f: impl FnOnce(&EnsembleSpaceConfig) -> T) -> Result<T, MLError> {
|
|
ENSEMBLE_CONFIG.with(|cell| {
|
|
let borrow = cell.borrow();
|
|
match borrow.as_ref() {
|
|
Some(cfg) => Ok(f(cfg)),
|
|
None => Err(MLError::ConfigError {
|
|
reason: "EnsembleSpaceConfig not installed. Call config.install() first."
|
|
.to_string(),
|
|
}),
|
|
}
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Parameter space
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Combined parameter space for ensemble optimization.
|
|
///
|
|
/// Holds a flat vector of continuous values:
|
|
/// `[model_0_params..., model_1_params..., ..., weight_0, weight_1, ...]`
|
|
///
|
|
/// The last N values (where N = number of models) are raw ensemble weights
|
|
/// in `[0, 1]`. Use [`extract_weights()`](Self::extract_weights) to get
|
|
/// normalized weights that sum to 1.
|
|
#[derive(Debug, Clone)]
|
|
pub struct EnsembleParameterSpace {
|
|
/// Flat continuous parameter vector.
|
|
pub values: Vec<f64>,
|
|
}
|
|
|
|
impl EnsembleParameterSpace {
|
|
/// Extract normalized ensemble weights from the parameter vector.
|
|
///
|
|
/// Weights are the last N values, normalized to sum to 1.
|
|
/// If all raw weights are zero (or negative after clamping), returns equal weights.
|
|
pub fn extract_weights(&self, num_models: usize) -> Vec<f64> {
|
|
let weight_start = self.values.len().saturating_sub(num_models);
|
|
let raw_weights: Vec<f64> = (0..num_models)
|
|
.filter_map(|i| self.values.get(weight_start + i).copied())
|
|
.map(|w| w.max(0.0))
|
|
.collect();
|
|
let sum: f64 = raw_weights.iter().sum();
|
|
if sum > f64::EPSILON {
|
|
raw_weights.iter().map(|w| w / sum).collect()
|
|
} else {
|
|
// Fallback: equal weights
|
|
let equal = 1.0 / num_models.max(1) as f64;
|
|
vec![equal; num_models]
|
|
}
|
|
}
|
|
|
|
/// Extract per-model parameter slices from the combined vector.
|
|
///
|
|
/// Returns a vector of slices, one per model. The slices are in the same
|
|
/// order as the model names in the config.
|
|
pub fn extract_model_params(&self, model_param_dims: &[usize]) -> Vec<Vec<f64>> {
|
|
let mut result = Vec::with_capacity(model_param_dims.len());
|
|
let mut offset = 0;
|
|
for &dim in model_param_dims {
|
|
let end = (offset + dim).min(self.values.len());
|
|
let start = offset.min(end);
|
|
result.push(self.values.get(start..end).unwrap_or_default().to_vec());
|
|
offset += dim;
|
|
}
|
|
result
|
|
}
|
|
}
|
|
|
|
impl ParameterSpace for EnsembleParameterSpace {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
// Read from thread-local config; if missing, return empty
|
|
// (caller should have installed config first)
|
|
with_config(|cfg| {
|
|
let mut bounds = cfg.model_bounds.clone();
|
|
// Append weight bounds: [0, 1] for each model
|
|
for _ in 0..cfg.num_models() {
|
|
bounds.push((0.0, 1.0));
|
|
}
|
|
bounds
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
let expected_dim = with_config(|cfg| cfg.total_dim())?;
|
|
if x.len() != expected_dim {
|
|
return Err(MLError::ConfigError {
|
|
reason: format!(
|
|
"EnsembleParameterSpace: expected {} params, got {}",
|
|
expected_dim,
|
|
x.len()
|
|
),
|
|
});
|
|
}
|
|
Ok(Self {
|
|
values: x.to_vec(),
|
|
})
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
self.values.clone()
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
// Build names from config. Since the trait requires &'static str,
|
|
// we leak the strings. This is acceptable because hyperopt configs
|
|
// are created once per optimization run (not in a hot loop).
|
|
with_config(|cfg| {
|
|
let mut names: Vec<&'static str> = cfg
|
|
.model_param_names
|
|
.iter()
|
|
.map(|s| -> &'static str { Box::leak(s.clone().into_boxed_str()) })
|
|
.collect();
|
|
for model_name in &cfg.model_names {
|
|
let weight_name = format!("weight_{}", model_name);
|
|
names.push(Box::leak(weight_name.into_boxed_str()));
|
|
}
|
|
names
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// Helper: create a simple 2-model config for testing.
|
|
fn test_config() -> EnsembleSpaceConfig {
|
|
EnsembleSpaceConfig::new(
|
|
vec!["model_a".into(), "model_b".into()],
|
|
vec![2, 3], // model_a has 2 params, model_b has 3 params
|
|
vec![
|
|
(0.0, 1.0),
|
|
(0.0, 10.0),
|
|
(-1.0, 1.0),
|
|
(-1.0, 1.0),
|
|
(0.0, 100.0),
|
|
],
|
|
vec![
|
|
"a_lr".into(),
|
|
"a_batch".into(),
|
|
"b_x".into(),
|
|
"b_y".into(),
|
|
"b_z".into(),
|
|
],
|
|
)
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_parameter_space_dimensions() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
// 2 + 3 model params + 2 weights = 7 total
|
|
assert_eq!(config.total_dim(), 7);
|
|
|
|
let bounds = EnsembleParameterSpace::continuous_bounds();
|
|
assert_eq!(bounds.len(), 7);
|
|
|
|
// First 5 bounds are model params
|
|
assert_eq!(bounds.first().copied(), Some((0.0, 1.0))); // a_lr
|
|
assert_eq!(bounds.get(4).copied(), Some((0.0, 100.0))); // b_z
|
|
|
|
// Last 2 bounds are weight params [0, 1]
|
|
assert_eq!(bounds.get(5).copied(), Some((0.0, 1.0))); // weight_model_a
|
|
assert_eq!(bounds.get(6).copied(), Some((0.0, 1.0))); // weight_model_b
|
|
|
|
let names = EnsembleParameterSpace::param_names();
|
|
assert_eq!(names.len(), 7);
|
|
assert_eq!(names.first().copied(), Some("a_lr"));
|
|
assert_eq!(names.get(4).copied(), Some("b_z"));
|
|
assert_eq!(names.get(5).copied(), Some("weight_model_a"));
|
|
assert_eq!(names.get(6).copied(), Some("weight_model_b"));
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_weights_normalizes() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
// 7 params: [a_lr, a_batch, b_x, b_y, b_z, weight_a, weight_b]
|
|
let params = EnsembleParameterSpace::from_continuous(&[
|
|
0.5, 5.0, 0.0, 0.0, 50.0, // model params
|
|
0.3, 0.7, // raw weights
|
|
])
|
|
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
|
|
|
let weights = params.extract_weights(2);
|
|
assert_eq!(weights.len(), 2);
|
|
|
|
// 0.3 / (0.3 + 0.7) = 0.3
|
|
assert!((weights.first().copied().unwrap_or(0.0) - 0.3).abs() < 1e-10);
|
|
// 0.7 / (0.3 + 0.7) = 0.7
|
|
assert!((weights.get(1).copied().unwrap_or(0.0) - 0.7).abs() < 1e-10);
|
|
|
|
// Verify weights sum to 1
|
|
let sum: f64 = weights.iter().sum();
|
|
assert!(
|
|
(sum - 1.0).abs() < 1e-10,
|
|
"Weights should sum to 1.0, got {}",
|
|
sum
|
|
);
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_weights_handles_zeros() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
let params = EnsembleParameterSpace::from_continuous(&[
|
|
0.5, 5.0, 0.0, 0.0, 50.0, // model params
|
|
0.0, 0.0, // all-zero weights
|
|
])
|
|
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
|
|
|
let weights = params.extract_weights(2);
|
|
assert_eq!(weights.len(), 2);
|
|
|
|
// Should fall back to equal weights: 0.5, 0.5
|
|
assert!(
|
|
(weights.first().copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
|
|
"Expected 0.5 for zero-weight fallback, got {:?}",
|
|
weights.first()
|
|
);
|
|
assert!(
|
|
(weights.get(1).copied().unwrap_or(0.0) - 0.5).abs() < 1e-10,
|
|
"Expected 0.5 for zero-weight fallback, got {:?}",
|
|
weights.get(1)
|
|
);
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_roundtrip_continuous() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
let original = vec![0.5, 5.0, 0.0, -0.5, 50.0, 0.3, 0.7];
|
|
let params = EnsembleParameterSpace::from_continuous(&original)
|
|
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
|
let recovered = params.to_continuous();
|
|
|
|
assert_eq!(original.len(), recovered.len());
|
|
for (a, b) in original.iter().zip(recovered.iter()) {
|
|
assert!(
|
|
(a - b).abs() < 1e-10,
|
|
"Round-trip mismatch: {} vs {}",
|
|
a,
|
|
b
|
|
);
|
|
}
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_wrong_dimension_returns_error() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
let result = EnsembleParameterSpace::from_continuous(&[1.0, 2.0]);
|
|
assert!(result.is_err(), "Should reject wrong dimension");
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_model_params() {
|
|
let config = test_config();
|
|
config.install();
|
|
|
|
let params = EnsembleParameterSpace::from_continuous(&[
|
|
0.5, 5.0, // model_a params (dim=2)
|
|
0.1, -0.3, 75.0, // model_b params (dim=3)
|
|
0.4, 0.6, // weights
|
|
])
|
|
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
|
|
|
let model_params = params.extract_model_params(&config.model_param_dims);
|
|
assert_eq!(model_params.len(), 2);
|
|
|
|
// model_a: [0.5, 5.0]
|
|
assert_eq!(model_params.first().map(|v| v.len()), Some(2));
|
|
assert!(
|
|
(model_params
|
|
.first()
|
|
.and_then(|v| v.first().copied())
|
|
.unwrap_or(0.0)
|
|
- 0.5)
|
|
.abs()
|
|
< 1e-10
|
|
);
|
|
|
|
// model_b: [0.1, -0.3, 75.0]
|
|
assert_eq!(model_params.get(1).map(|v| v.len()), Some(3));
|
|
assert!(
|
|
(model_params
|
|
.get(1)
|
|
.and_then(|v| v.get(2).copied())
|
|
.unwrap_or(0.0)
|
|
- 75.0)
|
|
.abs()
|
|
< 1e-10
|
|
);
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
|
|
#[test]
|
|
fn test_no_config_installed_returns_empty_or_error() {
|
|
// Ensure no config is installed
|
|
EnsembleSpaceConfig::uninstall();
|
|
|
|
// continuous_bounds returns empty when no config
|
|
let bounds = EnsembleParameterSpace::continuous_bounds();
|
|
assert!(bounds.is_empty());
|
|
|
|
// from_continuous returns error when no config
|
|
let result = EnsembleParameterSpace::from_continuous(&[1.0]);
|
|
assert!(result.is_err());
|
|
|
|
// param_names returns empty when no config
|
|
let names = EnsembleParameterSpace::param_names();
|
|
assert!(names.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_single_model_ensemble() {
|
|
let config = EnsembleSpaceConfig::new(
|
|
vec!["only_model".into()],
|
|
vec![1],
|
|
vec![(0.0, 1.0)],
|
|
vec!["lr".into()],
|
|
);
|
|
config.install();
|
|
|
|
assert_eq!(config.total_dim(), 2); // 1 param + 1 weight
|
|
|
|
let params = EnsembleParameterSpace::from_continuous(&[0.5, 0.8])
|
|
.unwrap_or_else(|e| panic!("from_continuous failed: {}", e));
|
|
let weights = params.extract_weights(1);
|
|
assert_eq!(weights.len(), 1);
|
|
// Single model: weight normalizes to 1.0
|
|
assert!(
|
|
(weights.first().copied().unwrap_or(0.0) - 1.0).abs() < 1e-10,
|
|
"Single model weight should be 1.0"
|
|
);
|
|
|
|
EnsembleSpaceConfig::uninstall();
|
|
}
|
|
}
|