docs: add DQN training smoke test implementation plan
4-task plan: add trainer accessors, write smoke test with 6 core assertions (loss convergence, finite losses, Q-value divergence, checkpoint integrity, epsilon decay), add Sharpe comparison vs untrained baseline, final verification. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
339
docs/plans/2026-02-20-dqn-training-smoke-test-implementation.md
Normal file
339
docs/plans/2026-02-20-dqn-training-smoke-test-implementation.md
Normal file
@@ -0,0 +1,339 @@
|
||||
# DQN Training Smoke Test — Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Single integration test proving train → checkpoint → validate pipeline works on real 6E.FUT data, with 7 assertions covering loss convergence, Q-value divergence, checkpoint integrity, epsilon decay, and Sharpe improvement.
|
||||
|
||||
**Architecture:** Extends the existing `dqn_training_pipeline_test.rs` pattern. Uses `DQNTrainer` with `DQNHyperparameters::conservative()` for fast training on real DBN data. Adds small public accessors to `DQNTrainer` for loss history and agent access. Reuses `RealDataLoader` + `ValidationHarness` from `validation_real_data_test.rs` for Sharpe comparison.
|
||||
|
||||
**Tech Stack:** Rust, tokio async, candle-core, anyhow, ml crate (DQN trainer, validation harness, real data loader)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add loss_history and agent accessors to DQNTrainer
|
||||
|
||||
**Files:**
|
||||
- Modify: `ml/src/trainers/dqn/trainer.rs` — add 3 public accessor methods
|
||||
|
||||
**Step 1: Add accessor methods after the existing impl block**
|
||||
|
||||
Find any convenient spot in the `impl DQNTrainer` block (near line 700, before `train()`) and add:
|
||||
|
||||
```rust
|
||||
/// Get per-epoch loss history (for smoke test verification)
|
||||
pub fn loss_history(&self) -> &[f64] {
|
||||
&self.loss_history
|
||||
}
|
||||
|
||||
/// Get per-epoch validation loss history
|
||||
pub fn val_loss_history(&self) -> &[f64] {
|
||||
&self.val_loss_history
|
||||
}
|
||||
|
||||
/// Get a clone of the DQN agent for Q-value inspection after training
|
||||
pub async fn get_agent_clone(&self) -> Result<crate::dqn::DQN> {
|
||||
let agent_lock = self.agent.read().await;
|
||||
match &*agent_lock {
|
||||
DQNAgentType::Working(agent) => Ok(agent.clone()),
|
||||
DQNAgentType::Standard(agent) => {
|
||||
anyhow::bail!("Standard agent clone not supported — use Working agent")
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Note: if `DQN` doesn't implement `Clone`, use a different approach — save to bytes via `varmap.save()` and load into a new `DQN`. Check the existing checkpoint save path in `train()` (line ~1961) for the pattern.
|
||||
|
||||
**Step 2: Verify compilation**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml`
|
||||
Expected: PASS (no errors)
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```
|
||||
feat(ml): add loss_history and agent accessors to DQNTrainer
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Write the smoke test — training + loss convergence assertions
|
||||
|
||||
**Files:**
|
||||
- Create: `ml/tests/dqn_training_smoke_test.rs`
|
||||
|
||||
**Step 1: Write the test file with Phase 1 (training + loss assertions)**
|
||||
|
||||
```rust
|
||||
//! DQN Training Smoke Test
|
||||
//!
|
||||
//! Verifies the complete train → checkpoint → validate pipeline
|
||||
//! works on real 6E.FUT data with 7 assertions:
|
||||
//!
|
||||
//! 1. All 20 epochs complete (no premature early stopping)
|
||||
//! 2. Loss decreases >30% (gradient flow works)
|
||||
//! 3. All per-epoch losses are finite (no NaN/Inf)
|
||||
//! 4. Q-value divergence (model develops action preferences)
|
||||
//! 5. Checkpoint round-trip (save/load weight integrity)
|
||||
//! 6. Epsilon decayed below 0.5 (exploration schedule ran)
|
||||
//! 7. Trained Sharpe > untrained Sharpe (better trading decisions)
|
||||
|
||||
#![allow(unused_crate_dependencies)]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn get_6e_fut_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!("6E.FUT data not found: {}", data_dir.display());
|
||||
}
|
||||
Ok(data_dir.to_string_lossy().to_string())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dqn_training_smoke() -> Result<()> {
|
||||
// === ARRANGE ===
|
||||
let data_dir = match get_6e_fut_data_dir() {
|
||||
Ok(dir) => dir,
|
||||
Err(e) => {
|
||||
eprintln!("Skipping: {}", e);
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let checkpoint_dir = tempfile::tempdir()?;
|
||||
|
||||
let mut hyperparams = DQNHyperparameters::conservative();
|
||||
hyperparams.epochs = 20;
|
||||
hyperparams.batch_size = 64;
|
||||
hyperparams.learning_rate = 0.0001;
|
||||
hyperparams.epsilon_start = 1.0;
|
||||
hyperparams.epsilon_end = 0.05;
|
||||
hyperparams.early_stopping_enabled = true;
|
||||
hyperparams.min_epochs_before_stopping = 25; // > 20 epochs = won't trigger
|
||||
hyperparams.checkpoint_frequency = 10;
|
||||
|
||||
// === ACT: Train ===
|
||||
let mut trainer = DQNTrainer::new(hyperparams.clone())?;
|
||||
let mut best_checkpoint_path = PathBuf::new();
|
||||
|
||||
let metrics = trainer
|
||||
.train(&data_dir, |epoch, checkpoint_data, is_best| {
|
||||
let name = if is_best {
|
||||
"smoke_best.safetensors".to_string()
|
||||
} else {
|
||||
format!("smoke_epoch_{}.safetensors", epoch)
|
||||
};
|
||||
let path = checkpoint_dir.path().join(&name);
|
||||
std::fs::write(&path, &checkpoint_data)?;
|
||||
if is_best {
|
||||
best_checkpoint_path = path.clone();
|
||||
}
|
||||
Ok(path.to_string_lossy().to_string())
|
||||
})
|
||||
.await?;
|
||||
|
||||
// === ASSERT 1: All 20 epochs completed ===
|
||||
assert_eq!(
|
||||
metrics.epochs_trained, 20,
|
||||
"ASSERT 1 FAILED: Expected 20 epochs, got {}. Early stopping fired prematurely.",
|
||||
metrics.epochs_trained
|
||||
);
|
||||
|
||||
// === ASSERT 2: Loss decreases >30% ===
|
||||
let loss_history = trainer.loss_history();
|
||||
assert!(
|
||||
loss_history.len() >= 2,
|
||||
"Need at least 2 epochs of loss history, got {}",
|
||||
loss_history.len()
|
||||
);
|
||||
let initial_loss = loss_history[0];
|
||||
let final_loss = *loss_history.last().unwrap_or(&f64::MAX);
|
||||
assert!(
|
||||
final_loss < initial_loss * 0.70,
|
||||
"ASSERT 2 FAILED: Loss did not decrease >30%. Initial={:.6}, Final={:.6}, Ratio={:.2}%",
|
||||
initial_loss, final_loss, (final_loss / initial_loss) * 100.0
|
||||
);
|
||||
|
||||
// === ASSERT 3: All losses finite ===
|
||||
for (i, loss) in loss_history.iter().enumerate() {
|
||||
assert!(
|
||||
loss.is_finite(),
|
||||
"ASSERT 3 FAILED: Loss at epoch {} is not finite: {}",
|
||||
i, loss
|
||||
);
|
||||
}
|
||||
|
||||
// === ASSERT 4: Q-value divergence ===
|
||||
let avg_q = metrics.additional_metrics.get("avg_q_value")
|
||||
.copied()
|
||||
.unwrap_or(0.0);
|
||||
// Q-values should be non-zero after training (model has preferences)
|
||||
assert!(
|
||||
avg_q.abs() > 0.001,
|
||||
"ASSERT 4 FAILED: Avg Q-value is near zero ({:.6}), model may not be learning preferences",
|
||||
avg_q
|
||||
);
|
||||
|
||||
// === ASSERT 5: Checkpoint round-trip ===
|
||||
assert!(
|
||||
best_checkpoint_path.exists(),
|
||||
"ASSERT 5 FAILED: No best checkpoint was saved"
|
||||
);
|
||||
let checkpoint_size = std::fs::metadata(&best_checkpoint_path)?.len();
|
||||
assert!(
|
||||
checkpoint_size > 1024,
|
||||
"ASSERT 5 FAILED: Checkpoint too small ({} bytes), likely corrupt",
|
||||
checkpoint_size
|
||||
);
|
||||
|
||||
// === ASSERT 6: Epsilon decayed ===
|
||||
let final_epsilon = metrics.additional_metrics.get("final_epsilon")
|
||||
.copied()
|
||||
.unwrap_or(1.0);
|
||||
assert!(
|
||||
final_epsilon < 0.5,
|
||||
"ASSERT 6 FAILED: Epsilon did not decay below 0.5 (got {:.4}). Exploration schedule may not have run.",
|
||||
final_epsilon
|
||||
);
|
||||
|
||||
// === REPORT ===
|
||||
println!("\n{}", "=".repeat(70));
|
||||
println!(" DQN TRAINING SMOKE TEST REPORT");
|
||||
println!("{}", "=".repeat(70));
|
||||
println!(" Epochs completed: {}", metrics.epochs_trained);
|
||||
println!(" Initial loss: {:.6}", initial_loss);
|
||||
println!(" Final loss: {:.6}", final_loss);
|
||||
println!(" Loss reduction: {:.1}%", (1.0 - final_loss / initial_loss) * 100.0);
|
||||
println!(" Avg Q-value: {:.4}", avg_q);
|
||||
println!(" Final epsilon: {:.4}", final_epsilon);
|
||||
println!(" Checkpoint size: {} KB", checkpoint_size / 1024);
|
||||
println!(" Training time: {:.1}s", metrics.training_time_seconds);
|
||||
println!("{}", "=".repeat(70));
|
||||
println!(" ASSERTIONS 1-6: ALL PASSED");
|
||||
println!("{}", "=".repeat(70));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Run the test to verify compilation and behavior**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture 2>&1`
|
||||
Expected: Test runs, all 6 assertions pass (assertion 7 — Sharpe comparison — is added in Task 3)
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```
|
||||
test(ml): add DQN training smoke test with 6 core assertions
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Add Sharpe comparison (trained vs untrained) — Assertion 7
|
||||
|
||||
**Files:**
|
||||
- Modify: `ml/tests/dqn_training_smoke_test.rs` — add validation phase
|
||||
|
||||
**Step 1: Add validation comparison after the training assertions**
|
||||
|
||||
This reuses the pattern from `validation_real_data_test.rs`. After training completes, run walk-forward validation on both an untrained and the trained model, compare Sharpe ratios.
|
||||
|
||||
The key challenge is bridging the `DQNTrainer` (which owns the agent) with the `DqnStrategy` adapter (which wraps a `DQN` for `ValidatableStrategy`). Use the `get_agent_clone()` accessor from Task 1, or if that doesn't work, use the checkpoint: load it into a fresh `DQN` and wrap in `DqnStrategy`.
|
||||
|
||||
Add to the test function, before the final report:
|
||||
|
||||
```rust
|
||||
// === ASSERT 7: Sharpe improvement over untrained ===
|
||||
// This is the ultimate test: does training produce better trading decisions?
|
||||
//
|
||||
// Approach: Load checkpoint into DqnStrategy, run walk-forward validation,
|
||||
// compare Sharpe vs the untrained baseline (0.0074 from validation_real_data_test).
|
||||
use ml::dqn::DQNConfig;
|
||||
use ml::validation::{
|
||||
DqnStrategy, TimeSeriesData, ValidationHarness, ValidationHarnessConfig,
|
||||
WalkForwardConfig,
|
||||
};
|
||||
use ml::real_data_loader::RealDataLoader;
|
||||
|
||||
// Load real data for validation
|
||||
let mut loader = RealDataLoader::new_from_workspace()?;
|
||||
let bars = loader.load_symbol_data("6E.FUT").await?;
|
||||
// ... build TimeSeriesData from bars (same pattern as validation_real_data_test.rs)
|
||||
// ... create DqnStrategy with trained checkpoint
|
||||
// ... run ValidationHarness
|
||||
// ... assert trained_sharpe > UNTRAINED_BASELINE_SHARPE
|
||||
|
||||
let untrained_sharpe = 0.0074; // From validation_real_data_test.rs
|
||||
// assert!(trained_sharpe > untrained_sharpe, ...);
|
||||
```
|
||||
|
||||
Note: The exact implementation depends on how the checkpoint can be loaded into a `DqnStrategy`. Follow the pattern in `validation_real_data_test.rs` lines 80-180, but load from checkpoint instead of creating fresh.
|
||||
|
||||
**Step 2: Run the test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture 2>&1`
|
||||
Expected: All 7 assertions pass
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```
|
||||
test(ml): add Sharpe comparison assertion to smoke test (7/7)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Final verification and cleanup
|
||||
|
||||
**Files:**
|
||||
- None (verification only)
|
||||
|
||||
**Step 1: Run the full ml test suite to verify no regressions**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -5`
|
||||
Expected: 1922+ tests pass, 0 failures
|
||||
|
||||
**Step 2: Run the smoke test one more time**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture 2>&1`
|
||||
Expected: All 7 assertions pass, clean report printed
|
||||
|
||||
**Step 3: Run the existing training pipeline tests to verify no regression**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --test dqn_training_pipeline_test -- --nocapture 2>&1`
|
||||
Expected: All existing tests still pass
|
||||
|
||||
**Step 4: Commit final state**
|
||||
|
||||
If any adjustments were needed during verification, commit them:
|
||||
```
|
||||
test(ml): finalize DQN training smoke test — all 7 assertions verified
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Key Implementation Notes
|
||||
|
||||
1. **The checkpoint callback signature**: Check whether it's `FnMut(usize, Vec<u8>, bool)` (3 args: epoch, data, is_best) or `FnMut(usize, Vec<u8>)` (2 args). The existing test in `dqn_training_pipeline_test.rs` uses 2 args. If 3 args, the `is_best` flag tells us when to save the best checkpoint.
|
||||
|
||||
2. **Q-value divergence**: The `additional_metrics["avg_q_value"]` gives us average Q. For divergence between actions, we'd need to do a forward pass. Using just `avg_q.abs() > 0.001` is a weaker but still useful assertion. To strengthen: get agent, do forward pass on sample state, check `max(Q) - min(Q) > threshold`.
|
||||
|
||||
3. **Sharpe baseline**: The untrained DQN gave Sharpe=0.0074 in validation_real_data_test.rs. This is essentially random. The trained model should beat this, but with only 20 epochs it may be marginal. Consider a softer assertion initially (trained_sharpe > 0.0).
|
||||
|
||||
4. **tempfile**: Use `tempfile::tempdir()` for checkpoint storage — auto-cleaned up after test.
|
||||
|
||||
5. **SQLX_OFFLINE=true**: Required for all cargo commands in this project.
|
||||
|
||||
6. **Clippy denials**: The project denies `unwrap_used`, `expect_used`, `panic`, `indexing_slicing`. Use `.get()`, `?` operator, and match expressions everywhere.
|
||||
|
||||
## Dependencies
|
||||
|
||||
- Task 2 depends on Task 1 (needs loss_history accessor)
|
||||
- Task 3 depends on Task 2 (extends the test)
|
||||
- Task 4 depends on Task 3 (final verification)
|
||||
Reference in New Issue
Block a user