1ad57f1bd5
After training completes, the example now loads the best checkpoint into a fresh DQN and runs inference on 5 synthetic state vectors, printing action, max Q-value, and Q-spread for each sample. Errors are handled gracefully with match on Result so the example never panics on inference failure. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
291 lines
10 KiB
Rust
291 lines
10 KiB
Rust
//! Production DQN Training Script
|
|
//!
|
|
//! Trains a DQN model for 50 epochs using production hyperparameters.
|
|
|
|
use anyhow::{Context, Result};
|
|
use candle_core::Tensor;
|
|
use ml::dqn::{DQNConfig, DQN};
|
|
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
|
use std::path::PathBuf;
|
|
use std::time::Instant;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Setup logging
|
|
tracing_subscriber::fmt()
|
|
.with_max_level(tracing::Level::INFO)
|
|
.init();
|
|
|
|
println!("\n{}", "=".repeat(80));
|
|
println!("🚀 DQN Production Training - 50 Epochs");
|
|
println!("{}", "=".repeat(80));
|
|
|
|
let start_time = Instant::now();
|
|
|
|
// Get data directory
|
|
let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
|
.parent()
|
|
.context("Failed to get workspace root")?
|
|
.to_path_buf();
|
|
|
|
let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
|
|
|
|
if !data_dir.exists() {
|
|
anyhow::bail!(
|
|
"Data directory not found: {}. Please check the path.",
|
|
data_dir.display()
|
|
);
|
|
}
|
|
|
|
// Create checkpoint directory
|
|
let checkpoint_dir = PathBuf::from("/tmp");
|
|
std::fs::create_dir_all(&checkpoint_dir)?;
|
|
|
|
println!("\n📋 Configuration:");
|
|
println!(" Data Directory: {}", data_dir.display());
|
|
println!(" Checkpoint Directory: {}", checkpoint_dir.display());
|
|
|
|
// Configure production hyperparameters (conservative baseline)
|
|
let mut hyperparams = DQNHyperparameters::conservative();
|
|
hyperparams.epochs = 50;
|
|
hyperparams.batch_size = 64;
|
|
hyperparams.learning_rate = 0.0001;
|
|
hyperparams.gamma = 0.99;
|
|
hyperparams.epsilon_start = 0.3;
|
|
hyperparams.epsilon_end = 0.05;
|
|
hyperparams.epsilon_decay = 0.995;
|
|
hyperparams.checkpoint_frequency = 10;
|
|
hyperparams.early_stopping_enabled = true;
|
|
hyperparams.min_epochs_before_stopping = 50; // Allow all 50 epochs
|
|
|
|
println!("\n⚙️ Hyperparameters:");
|
|
println!(" Epochs: {}", hyperparams.epochs);
|
|
println!(" Batch Size: {}", hyperparams.batch_size);
|
|
println!(" Learning Rate: {}", hyperparams.learning_rate);
|
|
println!(" Gamma: {}", hyperparams.gamma);
|
|
println!(
|
|
" Epsilon: {} → {} (decay: {})",
|
|
hyperparams.epsilon_start, hyperparams.epsilon_end, hyperparams.epsilon_decay
|
|
);
|
|
|
|
// Create trainer
|
|
println!("\n🏗️ Initializing DQN trainer...");
|
|
let mut trainer = DQNTrainer::new(hyperparams.clone())?;
|
|
|
|
// Train the model
|
|
println!("\n🚀 Starting training...\n");
|
|
|
|
let mut best_checkpoint_path = PathBuf::new();
|
|
|
|
let metrics = trainer
|
|
.train(
|
|
&data_dir.to_string_lossy().to_string(),
|
|
|epoch, checkpoint_data, is_best| {
|
|
let filename = if is_best {
|
|
"dqn_prod_best.safetensors".to_string()
|
|
} else {
|
|
format!("dqn_prod_epoch_{}.safetensors", epoch)
|
|
};
|
|
let path = checkpoint_dir.join(filename);
|
|
std::fs::write(&path, checkpoint_data)?;
|
|
if is_best {
|
|
best_checkpoint_path = path.clone();
|
|
println!(
|
|
" 💾 ⭐ BEST checkpoint saved: epoch {} -> {}",
|
|
epoch,
|
|
path.display()
|
|
);
|
|
} else {
|
|
println!(
|
|
" 💾 Checkpoint saved: epoch {} -> {}",
|
|
epoch,
|
|
path.display()
|
|
);
|
|
}
|
|
Ok(path.to_string_lossy().to_string())
|
|
},
|
|
)
|
|
.await?;
|
|
|
|
let training_time = start_time.elapsed();
|
|
|
|
// Report results
|
|
println!("\n{}", "=".repeat(80));
|
|
println!("✅ TRAINING COMPLETE");
|
|
println!("{}", "=".repeat(80));
|
|
println!("\n📊 Results:");
|
|
println!(" Epochs Completed: {}", metrics.epochs_trained);
|
|
println!(" Final Loss: {:.6}", metrics.loss);
|
|
println!(
|
|
" Training Time: {:.2}s ({:.1} min)",
|
|
training_time.as_secs_f64(),
|
|
training_time.as_secs_f64() / 60.0
|
|
);
|
|
println!(" Convergence: {}", metrics.convergence_achieved);
|
|
|
|
if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
|
|
println!(" Avg Q-value: {:.4}", avg_q_value);
|
|
}
|
|
|
|
if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
|
|
println!(" Final Epsilon: {:.4}", final_epsilon);
|
|
}
|
|
|
|
println!("\n💾 Best Checkpoint: {}", best_checkpoint_path.display());
|
|
|
|
let checkpoint_size = std::fs::metadata(&best_checkpoint_path)?.len();
|
|
println!(" Size: {} KB", checkpoint_size / 1024);
|
|
|
|
// =====================================================================
|
|
// Inference Demo: load best checkpoint and run on synthetic states
|
|
// =====================================================================
|
|
println!("\n{}", "=".repeat(80));
|
|
println!("INFERENCE DEMO");
|
|
println!("{}", "=".repeat(80));
|
|
|
|
let inference_result: Result<()> = (|| -> Result<()> {
|
|
let checkpoint_str = best_checkpoint_path
|
|
.to_str()
|
|
.context("Non-UTF8 checkpoint path")?;
|
|
|
|
// Load checkpoint tensors to discover architecture params
|
|
let checkpoint_tensors =
|
|
candle_core::safetensors::load(checkpoint_str, &candle_core::Device::Cpu)?;
|
|
|
|
// Detect noisy nets from key names
|
|
let uses_noisy = checkpoint_tensors.keys().any(|k| k.starts_with("noisy_"));
|
|
|
|
// Discover state_dim from first layer weight tensor
|
|
let state_dim = if uses_noisy {
|
|
checkpoint_tensors
|
|
.iter()
|
|
.find(|(name, _)| name.contains("noisy_hidden_0") && name.contains("mu_w"))
|
|
.map(|(_, t)| {
|
|
let d = t.dims();
|
|
if d.len() == 2 { d[1] } else { 54 }
|
|
})
|
|
.unwrap_or(54)
|
|
} else {
|
|
checkpoint_tensors
|
|
.iter()
|
|
.find(|(name, _)| name.contains("hidden_0") && name.contains("weight"))
|
|
.map(|(_, t)| {
|
|
let d = t.dims();
|
|
if d.len() == 2 { d[1] } else { 54 }
|
|
})
|
|
.unwrap_or(54)
|
|
};
|
|
|
|
// Build matching config
|
|
let mut config = DQNConfig::conservative();
|
|
config.state_dim = state_dim;
|
|
config.num_actions = 45;
|
|
config.hidden_dims = vec![256, 128, 64];
|
|
config.use_noisy_nets = uses_noisy;
|
|
config.noisy_sigma_init = 0.5;
|
|
config.use_iqn = true;
|
|
config.use_cql = true;
|
|
|
|
let num_actions = config.num_actions;
|
|
|
|
println!(
|
|
" Config: state_dim={}, num_actions={}, noisy={}",
|
|
state_dim, num_actions, uses_noisy
|
|
);
|
|
|
|
// Load into fresh DQN
|
|
let mut fresh_dqn = DQN::new(config)?;
|
|
fresh_dqn.load_from_safetensors(checkpoint_str)?;
|
|
let device = fresh_dqn.device().clone();
|
|
|
|
println!(" Loaded checkpoint into fresh DQN\n");
|
|
|
|
// Run inference on 5 synthetic state vectors
|
|
let num_samples: usize = 5;
|
|
for i in 0..num_samples {
|
|
// Deterministic synthetic state in [-0.5, 0.5]
|
|
let state_vec: Vec<f32> = (0..state_dim)
|
|
.map(|j| ((i * state_dim + j) as f32 * 0.037).sin() * 0.5)
|
|
.collect();
|
|
|
|
let state_tensor = match Tensor::from_vec(state_vec, (1, state_dim), &device) {
|
|
Ok(t) => t,
|
|
Err(e) => {
|
|
eprintln!(" [WARN] Sample {}: failed to create tensor: {}", i + 1, e);
|
|
continue;
|
|
}
|
|
};
|
|
|
|
match fresh_dqn.forward(&state_tensor) {
|
|
Ok(q_values) => {
|
|
let q_vec: Vec<f32> = q_values
|
|
.to_vec2::<f32>()?
|
|
.into_iter()
|
|
.flatten()
|
|
.collect();
|
|
|
|
// Find best action (argmax)
|
|
let (best_action, max_q) = q_vec
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, a), (_, b)| {
|
|
a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
|
|
})
|
|
.map(|(idx, &val)| (idx, val))
|
|
.unwrap_or((0, 0.0));
|
|
|
|
let min_q = q_vec
|
|
.iter()
|
|
.copied()
|
|
.min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
|
.unwrap_or(0.0);
|
|
|
|
let q_spread = max_q - min_q;
|
|
|
|
println!(
|
|
" Sample {}/{}: action={:>2} max_Q={:>+10.4} Q-spread={:.4}",
|
|
i + 1,
|
|
num_samples,
|
|
best_action,
|
|
max_q,
|
|
q_spread,
|
|
);
|
|
}
|
|
Err(e) => {
|
|
eprintln!(" [WARN] Sample {}: inference failed: {}", i + 1, e);
|
|
}
|
|
}
|
|
}
|
|
|
|
println!("\n Inference demo complete.");
|
|
Ok(())
|
|
})();
|
|
|
|
if let Err(e) = inference_result {
|
|
eprintln!(
|
|
"\n [WARN] Inference demo failed (training results above are still valid): {}",
|
|
e
|
|
);
|
|
}
|
|
|
|
println!("\n{}", "=".repeat(80));
|
|
|
|
// Save metrics to JSON
|
|
let metrics_json = serde_json::json!({
|
|
"epochs_trained": metrics.epochs_trained,
|
|
"final_loss": metrics.loss,
|
|
"training_time_seconds": training_time.as_secs_f64(),
|
|
"convergence_achieved": metrics.convergence_achieved,
|
|
"avg_q_value": metrics.additional_metrics.get("avg_q_value"),
|
|
"final_epsilon": metrics.additional_metrics.get("final_epsilon"),
|
|
"checkpoint_path": best_checkpoint_path.to_string_lossy().to_string(),
|
|
"checkpoint_size_kb": checkpoint_size / 1024,
|
|
});
|
|
|
|
let metrics_path = PathBuf::from("/tmp/dqn_production_test_training.json");
|
|
std::fs::write(&metrics_path, serde_json::to_string_pretty(&metrics_json)?)?;
|
|
println!("📄 Training metrics saved to: {}", metrics_path.display());
|
|
|
|
Ok(())
|
|
}
|