e2f6fc17f5
Validates the complete DQN checkpoint lifecycle: train 5 epochs with DQNHyperparameters::conservative(), save via checkpoint callback, load into a fresh DQN with architecture auto-detected from checkpoint tensor metadata (noisy vs standard layers, state_dim), and run 100 inference passes asserting valid action indices, finite Q-values, and non-zero Q-values. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
275 lines
10 KiB
Rust
275 lines
10 KiB
Rust
//! DQN Checkpoint -> Inference Integration Test
|
|
//!
|
|
//! Proves the complete checkpoint -> load -> inference path:
|
|
//! 1. Train 5 epochs on real market data (fast, just to get a valid checkpoint)
|
|
//! 2. Save checkpoint to temp dir via the trainer's checkpoint callback
|
|
//! 3. Load checkpoint into a fresh DQN with matching architecture
|
|
//! 4. Run inference on 100 synthetic state vectors
|
|
//! 5. Assert: actions valid (0..45), Q-values finite, Q-values non-zero
|
|
|
|
#![allow(unused_crate_dependencies)]
|
|
|
|
use anyhow::{Context, Result};
|
|
use candle_core::Tensor;
|
|
use ml::dqn::{DQNConfig, DQN};
|
|
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
|
use std::path::PathBuf;
|
|
|
|
/// Locate the small training data directory, returning an error if absent.
|
|
fn get_data_dir() -> Result<String> {
|
|
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 not found: {}", data_dir.display());
|
|
}
|
|
Ok(data_dir.to_string_lossy().to_string())
|
|
}
|
|
|
|
/// Build a `DQNConfig` that matches the trainer's internal config exactly.
|
|
///
|
|
/// The trainer builds its DQNConfig from DQNHyperparameters::conservative() with
|
|
/// hardcoded architecture params: state_dim=54, num_actions=45, hidden_dims=[256,128,64].
|
|
/// When `use_noisy_nets=true` (conservative default), the Sequential network uses
|
|
/// NoisyLinear layers with key prefix "noisy_hidden_*" / "noisy_output".
|
|
/// We detect the naming scheme from the checkpoint to guarantee a match.
|
|
fn build_matching_config(checkpoint_tensors: &std::collections::HashMap<String, Tensor>) -> DQNConfig {
|
|
// Detect whether the checkpoint was produced with noisy nets by inspecting
|
|
// the VarMap key prefixes.
|
|
let uses_noisy = checkpoint_tensors.keys().any(|k| k.starts_with("noisy_"));
|
|
|
|
// Discover state_dim from the first layer weight.
|
|
// NoisyLinear: "noisy_hidden_0.mu_w" shape [hidden, input]
|
|
// Standard: "hidden_0.weight" shape [hidden, input]
|
|
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)
|
|
};
|
|
|
|
// Replicate the exact config the trainer constructs (see trainer.rs ~line 289).
|
|
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;
|
|
// Trainer hardcodes noisy_sigma_init = 0.5 via hyperparams
|
|
config.noisy_sigma_init = 0.5;
|
|
// Trainer hardcodes use_iqn = true, use_cql = true
|
|
config.use_iqn = true;
|
|
config.use_cql = true;
|
|
config
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_to_inference() -> Result<()> {
|
|
// -- Skip gracefully when real data is absent (CI environments) --------
|
|
let data_dir = match get_data_dir() {
|
|
Ok(dir) => dir,
|
|
Err(e) => {
|
|
eprintln!("Skipping test_checkpoint_to_inference: {e}");
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
let checkpoint_dir = tempfile::tempdir()?;
|
|
|
|
// =====================================================================
|
|
// Phase 1: Train for 5 epochs to produce a valid checkpoint
|
|
// =====================================================================
|
|
let mut hyperparams = DQNHyperparameters::conservative();
|
|
hyperparams.epochs = 5;
|
|
hyperparams.batch_size = 32;
|
|
hyperparams.learning_rate = 0.0001;
|
|
hyperparams.early_stopping_enabled = false;
|
|
hyperparams.checkpoint_frequency = 5; // checkpoint on last epoch
|
|
|
|
let mut trainer = DQNTrainer::new(hyperparams)?;
|
|
let mut best_checkpoint_path: Option<PathBuf> = None;
|
|
|
|
let _metrics = trainer
|
|
.train(&data_dir, |epoch, checkpoint_data, is_best| {
|
|
let name = if is_best {
|
|
"inference_best.safetensors".to_string()
|
|
} else {
|
|
format!("inference_epoch_{epoch}.safetensors")
|
|
};
|
|
let path = checkpoint_dir.path().join(&name);
|
|
std::fs::write(&path, &checkpoint_data)?;
|
|
if is_best {
|
|
best_checkpoint_path = Some(path.clone());
|
|
}
|
|
Ok(path.to_string_lossy().to_string())
|
|
})
|
|
.await?;
|
|
|
|
// If no "best" was saved, fall back to any checkpoint in the directory.
|
|
let checkpoint_path = match best_checkpoint_path {
|
|
Some(p) => p,
|
|
None => {
|
|
// Find first .safetensors file in the temp dir
|
|
let mut entries: Vec<_> = std::fs::read_dir(checkpoint_dir.path())?
|
|
.filter_map(|e| e.ok())
|
|
.filter(|e| {
|
|
e.path()
|
|
.extension()
|
|
.map_or(false, |ext| ext == "safetensors")
|
|
})
|
|
.collect();
|
|
anyhow::ensure!(
|
|
!entries.is_empty(),
|
|
"No checkpoint files were saved during training"
|
|
);
|
|
entries.sort_by_key(|e| e.path());
|
|
entries.pop().context("No checkpoint file")?.path()
|
|
}
|
|
};
|
|
|
|
assert!(
|
|
checkpoint_path.exists(),
|
|
"Checkpoint file does not exist: {}",
|
|
checkpoint_path.display()
|
|
);
|
|
println!(
|
|
"Phase 1 complete: checkpoint at {} ({} bytes)",
|
|
checkpoint_path.display(),
|
|
std::fs::metadata(&checkpoint_path)?.len()
|
|
);
|
|
|
|
// =====================================================================
|
|
// Phase 2: Load checkpoint into a fresh DQN
|
|
// =====================================================================
|
|
// Load the checkpoint's tensor metadata to discover architecture params
|
|
// (state_dim, noisy vs standard layers) so we can construct a DQN whose
|
|
// VarMap key names and tensor shapes match exactly.
|
|
let checkpoint_path_str = checkpoint_path
|
|
.to_str()
|
|
.context("Non-UTF8 checkpoint path")?;
|
|
|
|
let checkpoint_tensors =
|
|
candle_core::safetensors::load(checkpoint_path_str, &candle_core::Device::Cpu)?;
|
|
|
|
println!(
|
|
"Checkpoint contains {} tensors: {:?}",
|
|
checkpoint_tensors.len(),
|
|
checkpoint_tensors.keys().collect::<Vec<_>>()
|
|
);
|
|
|
|
let config = build_matching_config(&checkpoint_tensors);
|
|
let state_dim = config.state_dim;
|
|
println!(
|
|
"Built matching config: state_dim={}, num_actions={}, noisy={}",
|
|
config.state_dim, config.num_actions, config.use_noisy_nets
|
|
);
|
|
|
|
let mut fresh_dqn = DQN::new(config)?;
|
|
fresh_dqn.load_from_safetensors(checkpoint_path_str)?;
|
|
|
|
println!("Phase 2 complete: fresh DQN loaded from checkpoint");
|
|
|
|
// =====================================================================
|
|
// Phase 3: Inference on 100 synthetic state vectors
|
|
// =====================================================================
|
|
let num_inference_samples: usize = 100;
|
|
let num_actions: usize = 45;
|
|
let device = fresh_dqn.device().clone();
|
|
|
|
let mut all_actions_valid = true;
|
|
let mut all_q_finite = true;
|
|
let mut any_q_nonzero = false;
|
|
let mut action_counts = vec![0usize; num_actions];
|
|
|
|
for i in 0..num_inference_samples {
|
|
// Deterministic synthetic state: slight variation per sample
|
|
let state_vec: Vec<f32> = (0..state_dim)
|
|
.map(|j| ((i * state_dim + j) as f32 * 0.01).sin())
|
|
.collect();
|
|
|
|
let state_tensor = Tensor::from_vec(state_vec, (1, state_dim), &device)?;
|
|
|
|
let q_values = fresh_dqn.forward(&state_tensor)?;
|
|
|
|
// Shape check: [1, 45]
|
|
let dims = q_values.dims();
|
|
assert_eq!(
|
|
dims,
|
|
&[1, num_actions],
|
|
"Q-value tensor shape mismatch: expected [1, {}], got {:?}",
|
|
num_actions,
|
|
dims
|
|
);
|
|
|
|
let q_vec: Vec<f32> = q_values.to_vec2::<f32>()?.into_iter().flatten().collect();
|
|
|
|
// Check finite
|
|
for &q in &q_vec {
|
|
if !q.is_finite() {
|
|
all_q_finite = false;
|
|
}
|
|
if q.abs() > 1e-9 {
|
|
any_q_nonzero = true;
|
|
}
|
|
}
|
|
|
|
// Argmax to get action index
|
|
let best_action_idx = q_vec
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
|
.map(|(idx, _)| idx)
|
|
.unwrap_or(0);
|
|
|
|
if best_action_idx >= num_actions {
|
|
all_actions_valid = false;
|
|
} else if let Some(count) = action_counts.get_mut(best_action_idx) {
|
|
*count += 1;
|
|
}
|
|
}
|
|
|
|
// =====================================================================
|
|
// Phase 4: Assertions
|
|
// =====================================================================
|
|
assert!(all_q_finite, "Some Q-values were NaN or Inf");
|
|
assert!(
|
|
any_q_nonzero,
|
|
"All Q-values were zero -- model likely failed to load weights"
|
|
);
|
|
assert!(
|
|
all_actions_valid,
|
|
"Some action indices were out of range (expected 0..45)"
|
|
);
|
|
|
|
// Report action distribution
|
|
let unique_actions = action_counts.iter().filter(|&&c| c > 0).count();
|
|
println!(
|
|
"Phase 3 complete: {num_inference_samples} inference passes, \
|
|
{unique_actions} unique actions selected"
|
|
);
|
|
|
|
// Sanity: at least 1 unique action (degenerate model is still valid for this test)
|
|
assert!(
|
|
unique_actions >= 1,
|
|
"No actions were selected -- inference loop failure"
|
|
);
|
|
|
|
println!("All assertions passed: checkpoint -> load -> inference pipeline verified");
|
|
|
|
Ok(())
|
|
}
|