gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(dqn,ppo): verification swarm round 2 — gamma^n, checkpoint config, circuit breaker

Browse Source 
- IQN + C51 Bellman bootstrap: gamma → gamma^n for n-step returns
- DQNConfig::from_safetensors_file(): reconstruct config from checkpoint
  metadata instead of hardcoding dims/Rainbow flags in enhanced_ml.rs
- PPO update_value_only(): add circuit breaker guard (consistency with update())
- PPO eval: tensor-core alignment on hidden dims (must match PpoTrainer)
- enhanced_ml.rs: checkpoint-driven config loading with fallback

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-04 00:40:08 +01:00
parent cdb32cd8bf
commit 042c38fed2
4 changed files with 173 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml/examples/evaluate_baseline.rs
View File

@@ -530,11 +530,13 @@ fn evaluate_ppo_fold(
num_actions: args.num_actions,
policy_hidden_dims: {
let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128);
vec![base, base / 2]
let align = |x: usize| x.div_ceil(8) * 8; // tensor-core alignment (must match PpoTrainer)
vec![align(base), align(base / 2)]
},
value_hidden_dims: {
let base = hp_usize(hp, "hidden_dim_base").unwrap_or(128);
vec![base * 4, base * 3, base * 2, base, base / 2] // 5-layer to match trainer
let align = |x: usize| x.div_ceil(8) * 8; // tensor-core alignment (must match PpoTrainer)
vec![align(base * 4), align(base * 3), align(base * 2), align(base), align(base / 2)]
},
policy_learning_rate: 3e-4,
value_learning_rate: 1e-3,

112
crates/ml/src/dqn/dqn.rs
View File

@@ -320,9 +320,112 @@ impl DQNConfig {
meta.insert("dqn.use_distributional".to_owned(), self.use_distributional.to_string());
meta.insert("dqn.use_noisy_nets".to_owned(), self.use_noisy_nets.to_string());
meta.insert("dqn.use_iqn".to_owned(), self.use_iqn.to_string());
meta.insert("dqn.dueling_hidden_dim".to_owned(), self.dueling_hidden_dim.to_string());
meta.insert("dqn.num_atoms".to_owned(), self.num_atoms.to_string());
meta.insert("dqn.iqn_embedding_dim".to_owned(), self.iqn_embedding_dim.to_string());
meta.insert("dqn.iqn_num_quantiles".to_owned(), self.iqn_num_quantiles.to_string());
meta.insert("dqn.gamma".to_owned(), self.gamma.to_string());
meta
}
/// Load DQNConfig from a safetensors checkpoint file.
///
/// Reads architecture-critical fields (state_dim, num_actions, hidden_dims,
/// Rainbow flags) from the file's metadata. Non-architecture fields (LR,
/// epsilon, buffer capacity) use defaults since they don't affect inference.
pub fn from_safetensors_file(path: &std::path::Path) -> Result<Self, crate::MLError> {
let raw_bytes = std::fs::read(path).map_err(|e| {
crate::MLError::CheckpointError(format!("Failed to read safetensors: {}", e))
})?;
let (_header_size, st_metadata) =
safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
crate::MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
})?;
let meta = match st_metadata.metadata() {
Some(m) => m,
None => {
return Err(crate::MLError::CheckpointError(
"Safetensors checkpoint has no metadata — re-train to embed config".to_owned(),
));
}
};
Self::from_checkpoint_metadata(meta)
}
/// Reconstruct a DQNConfig from safetensors checkpoint metadata.
///
/// Reads architecture-critical fields (state_dim, num_actions, hidden_dims,
/// Rainbow flags) from the metadata map. Non-architecture fields (LR, epsilon,
/// buffer capacity) use defaults since they don't affect inference.
///
/// Returns `Err` if required fields are missing or malformed.
pub fn from_checkpoint_metadata(
metadata: &std::collections::HashMap<String, String>,
) -> Result<Self, crate::MLError> {
let get = |key: &str| -> Result<&str, crate::MLError> {
metadata.get(key).map(|s| s.as_str()).ok_or_else(|| {
crate::MLError::CheckpointError(format!(
"Checkpoint metadata missing '{}' — re-train to embed config", key
))
})
};
let state_dim: usize = get("dqn.state_dim")?
.parse()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.state_dim: {}", e)))?;
let num_actions: usize = get("dqn.num_actions")?
.parse()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.num_actions: {}", e)))?;
// Parse hidden_dims from debug format "[256, 256]"
let hidden_dims_str = get("dqn.hidden_dims")?;
let hidden_dims: Vec<usize> = hidden_dims_str
.trim_matches(|c| c == '[' || c == ']')
.split(',')
.filter(|s| !s.trim().is_empty())
.map(|s| s.trim().parse::<usize>())
.collect::<Result<Vec<_>, _>>()
.map_err(|e| crate::MLError::CheckpointError(format!("Bad dqn.hidden_dims: {}", e)))?;
let parse_bool = |key: &str| -> Result<bool, crate::MLError> {
get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
};
let parse_usize = |key: &str| -> Result<usize, crate::MLError> {
get(key)?.parse().map_err(|e| crate::MLError::CheckpointError(format!("Bad {}: {}", key, e)))
};
let use_dueling = parse_bool("dqn.use_dueling")?;
let use_distributional = parse_bool("dqn.use_distributional")?;
let use_noisy_nets = parse_bool("dqn.use_noisy_nets")?;
let use_iqn = parse_bool("dqn.use_iqn")?;
let dueling_hidden_dim = parse_usize("dqn.dueling_hidden_dim")?;
let num_atoms = parse_usize("dqn.num_atoms")?;
let iqn_embedding_dim = parse_usize("dqn.iqn_embedding_dim")?;
let iqn_num_quantiles = parse_usize("dqn.iqn_num_quantiles")?;
let gamma: f32 = metadata.get("dqn.gamma")
.and_then(|s| s.parse().ok())
.unwrap_or(0.95);
Ok(Self {
state_dim,
num_actions,
hidden_dims,
use_dueling,
dueling_hidden_dim,
use_distributional,
num_atoms,
use_noisy_nets,
use_iqn,
iqn_embedding_dim,
iqn_num_quantiles,
gamma,
..Self::default()
})
}
/// Validate that a checkpoint's metadata matches this config's architecture.
///
/// Returns `Ok(())` if the hash matches. Returns an error if metadata is
@@ -1904,8 +2007,9 @@ impl DQN {
.gather(&next_actions_for_gather, 1)?
.squeeze(1)?; // [batch, num_quantiles]
// Compute target quantiles: r + γ * (1 - done) * Z_target
let gamma = self.config.gamma;
// Compute target quantiles: r + γ^n * (1 - done) * Z_target
// Must use gamma^n for n-step returns (NStepBuffer accumulates n-step rewards)
let gamma = self.config.gamma.powi(self.config.n_steps as i32);
let rewards_broadcast = rewards_tensor
.unsqueeze(1)?
.broadcast_as((batch_size, num_quantiles))?;
@@ -2020,12 +2124,14 @@ impl DQN {
};
// Apply Bellman operator to get target distribution
// Must use gamma^n for n-step returns (NStepBuffer accumulates n-step rewards)
let gamma_n = self.config.gamma.powi(self.config.n_steps as i32);
let target_dists = if let Some(ref cat_dist) = self.categorical_dist {
cat_dist.apply_bellman_operator(
&rewards_tensor,
&next_dists,
&dones_tensor,
self.config.gamma,
gamma_n,
)?
} else {
return Err(MLError::TrainingError(

8
crates/ml/src/ppo/ppo.rs
View File

@@ -1029,6 +1029,14 @@ impl PPO {
///
/// Returns the average value loss across all mini-batches and epochs.
pub fn update_value_only(&mut self, batch: &mut TrajectoryBatch) -> Result<f32, MLError> {
// Circuit breaker guard (consistent with update())
if let Some(ref circuit_breaker) = self.circuit_breaker {
if !circuit_breaker.allow_request() {
warn!("Circuit breaker is open - skipping value-only update");
return Ok(0.0);
}
}
// Initialize optimizers (we only use value_optimizer, but init both for consistency)
self.init_optimizers()?;

77
services/trading_service/src/services/enhanced_ml.rs
View File

@@ -1444,31 +1444,7 @@ impl RealDQNModel {
use ml::dqn::agent::DQNAgent;
use ml::dqn::DQNConfig;
// TODO: Load architecture config from checkpoint metadata instead of hardcoding
// DQN configuration matching production training defaults
// (51 market + 3 portfolio features, 5 exposure x 3 order x 3 urgency factored actions)
let config = DQNConfig {
state_dim: 54,
num_actions: 45,
hidden_dims: vec![128, 64, 32],
learning_rate: 0.0001,
gamma: 0.99,
epsilon_start: 0.1,
epsilon_end: 0.01,
epsilon_decay: 0.995,
replay_buffer_capacity: 100000,
batch_size: 128,
target_update_freq: 1000,
minimum_profit_factor: 1.5,
tau: 0.005,
weight_decay: 1e-4,
..DQNConfig::default()
};
let mut agent = DQNAgent::new(config)
.map_err(|e| ml::MLError::ModelError(format!("Failed to create DQN agent: {}", e)))?;
// Try safetensors first, fall back to JSON checkpoint
// Resolve safetensors path
let safetensors_path = if checkpoint_path
.extension()
.is_some_and(|ext| ext == "safetensors")
@@ -1483,6 +1459,37 @@ impl RealDQNModel {
}
};
// Load architecture config from checkpoint metadata when available.
// This ensures hidden_dims, Rainbow flags, etc. match exactly what was
// trained, regardless of hyperopt-selected architecture.
let config = if let Some(ref st_path) = safetensors_path {
match DQNConfig::from_safetensors_file(st_path) {
Ok(mut cfg) => {
// Override inference-only settings
cfg.epsilon_start = 0.0; // No exploration during inference
cfg.epsilon_end = 0.0;
tracing::info!(
"DQN config loaded from checkpoint metadata: state_dim={}, num_actions={}, hidden_dims={:?}",
cfg.state_dim, cfg.num_actions, cfg.hidden_dims
);
cfg
}
Err(e) => {
tracing::warn!(
"Failed to read DQN config from checkpoint metadata, using defaults: {}",
e
);
Self::fallback_dqn_config()
}
}
} else {
// JSON checkpoint — no metadata available
Self::fallback_dqn_config()
};
let mut agent = DQNAgent::new(config)
.map_err(|e| ml::MLError::ModelError(format!("Failed to create DQN agent: {}", e)))?;
if let Some(st_path) = safetensors_path {
match agent.load_from_safetensors(&st_path) {
Ok(()) => {
@@ -1518,6 +1525,26 @@ impl RealDQNModel {
feature_count: 54,
})
}
/// Fallback DQN config for legacy JSON checkpoints without metadata.
fn fallback_dqn_config() -> ml::dqn::DQNConfig {
ml::dqn::DQNConfig {
state_dim: 54,
num_actions: 45,
hidden_dims: vec![128, 64, 32],
learning_rate: 0.0001,
gamma: 0.95,
epsilon_start: 0.0,
epsilon_end: 0.0,
replay_buffer_capacity: 100_000,
batch_size: 128,
target_update_freq: 1000,
minimum_profit_factor: 1.5,
tau: 0.005,
weight_decay: 1e-4,
..ml::dqn::DQNConfig::default()
}
}
}
#[async_trait::async_trait]