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

refactor(ml): delete straggler train_mamba2/train_ppo examples

Browse Source 
These two files survived the 20-file consolidation in 022036cb.
Both are now fully superseded:
- train_ppo.rs → train_baseline_rl --model ppo
- train_mamba2.rs → train_baseline_supervised --model mamba2

Also updates entrypoint-generic.sh usage examples to reference
the unified binaries (train_baseline_rl, train_baseline_supervised).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-25 23:24:53 +01:00
parent 267240530d
commit 71eefa6d53
3 changed files with 8 additions and 671 deletions
Download Patch File Download Diff File Expand all files Collapse all files

257
crates/ml/examples/train_mamba2.rs
View File

@@ -1,257 +0,0 @@
//! MAMBA-2 Training Example
//!
//! Trains a MAMBA-2 state space model on real market data from DBN files.
//! Uses continuous price sequences with microstructure features for next-timestep prediction.
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (100 epochs, real DBN data)
//! cargo run -p ml --example train_mamba2 --release --features cuda
//!
//! # Custom parameters with specific DBN directory
//! cargo run -p ml --example train_mamba2 --release --features cuda -- \
//! --epochs 500 \
//! --d-model 256 \
//! --n-layers 6 \
//! --seq-len 60 \
//! --dbn-dir test_data/real/databento/ml_training_small
//!
//! # Quick test with fewer epochs
//! cargo run -p ml --example train_mamba2 --release --features cuda -- \
//! --epochs 10 \
//! --batch-size 4
//! ```
//!
//! # Data Requirements
//!
//! - DBN files with OHLCV data (1-minute bars recommended)
//! - At least 60-128 consecutive timesteps per symbol
//! - Multiple files per symbol for better training data
use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tracing::info;
use tracing_subscriber::FmtSubscriber;
use ml::data_loaders::DbnSequenceLoader;
use ml::trainers::mamba2::{Mamba2Hyperparameters, Mamba2Trainer};
#[derive(Debug, Parser)]
#[command(name = "train_mamba2", about = "Train MAMBA-2 model on market data")]
struct Opts {
/// Number of training epochs
#[arg(long, default_value = "100")]
epochs: usize,
/// Learning rate
#[arg(long, default_value = "0.0001")]
learning_rate: f64,
/// Batch size (1-16 for 4GB VRAM)
#[arg(long, default_value = "8")]
batch_size: usize,
/// Model dimension (256, 512, 1024)
#[arg(long, default_value = "256")]
d_model: usize,
/// Number of layers (4-12)
#[arg(long, default_value = "6")]
n_layers: usize,
/// Sequence length
#[arg(long, default_value = "128")]
seq_len: usize,
/// Output directory for trained model
#[arg(long, default_value = "ml/trained_models")]
output_dir: String,
/// DBN data directory (contains .dbn files)
#[arg(long, default_value = "test_data/real/databento/ml_training_small")]
dbn_dir: String,
/// Train/validation split ratio
#[arg(long, default_value = "0.9")]
train_split: f64,
/// Verbose logging
#[arg(short, long)]
verbose: bool,
}
#[tokio::main]
async fn main() -> Result<()> {
// Parse CLI options
let opts = Opts::parse();
// Setup logging
let level = if opts.verbose {
tracing::Level::DEBUG
} else {
tracing::Level::INFO
};
let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
tracing::subscriber::set_global_default(subscriber)
.context("Failed to set tracing subscriber")?;
info!("🚀 Starting MAMBA-2 Training");
info!("Configuration:");
info!(" • Epochs: {}", opts.epochs);
info!(" • Learning rate: {}", opts.learning_rate);
info!(" • Batch size: {}", opts.batch_size);
info!(" • Model dimension: {}", opts.d_model);
info!(" • Number of layers: {}", opts.n_layers);
info!(" • Sequence length: {}", opts.seq_len);
info!(" • Output directory: {}", opts.output_dir);
// Create output directory
let output_path = PathBuf::from(&opts.output_dir);
if !output_path.exists() {
std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
info!("✅ Created output directory: {}", opts.output_dir);
}
// Configure MAMBA-2 hyperparameters
let hyperparams = Mamba2Hyperparameters {
learning_rate: opts.learning_rate,
batch_size: opts.batch_size,
d_model: opts.d_model,
n_layers: opts.n_layers,
state_size: 32,
dropout: 0.1,
epochs: opts.epochs,
seq_len: opts.seq_len,
grad_clip: 1.0,
weight_decay: 1e-4,
warmup_steps: 1000,
};
// Validate hyperparameters for VRAM constraint
hyperparams
.validate()
.context("Invalid hyperparameters for 4GB VRAM")?;
info!(
"✅ Hyperparameters validated (estimated VRAM: {}MB)",
hyperparams.estimate_memory_usage()
);
// Create MAMBA-2 trainer
let checkpoint_path = format!("{}/mamba2", opts.output_dir);
let mut trainer = Mamba2Trainer::new(hyperparams.clone(), Some(checkpoint_path))
.context("Failed to create MAMBA-2 trainer")?;
info!(
"✅ MAMBA-2 trainer initialized (job_id: {})",
trainer.job_id
);
// Load real DBN market data sequences
info!("\n📊 Loading DBN market data sequences...");
info!(" • DBN directory: {}", opts.dbn_dir);
info!(" • Sequence length: {}", opts.seq_len);
info!(" • Feature dimension: {}", opts.d_model);
info!(
" • Train/val split: {:.1}/{:.1}",
opts.train_split * 100.0,
(1.0 - opts.train_split) * 100.0
);
let mut loader = DbnSequenceLoader::new(opts.seq_len, opts.d_model)
.await
.context("Failed to create DBN sequence loader")?;
let (train_data, val_data) = loader
.load_sequences(&opts.dbn_dir, opts.train_split)
.await
.context("Failed to load DBN sequences")?;
info!(
"✅ Loaded {} training sequences, {} validation sequences",
train_data.len(),
val_data.len()
);
if train_data.is_empty() {
return Err(anyhow::anyhow!(
"No training sequences loaded! Check DBN directory: {}",
opts.dbn_dir
));
}
// Log sequence shape information
if let Some((input, target)) = train_data.first() {
info!(" • Input shape: {:?}", input.dims());
info!(" • Target shape: {:?}", target.dims());
}
// Set progress callback
let progress_callback =
std::sync::Arc::new(move |progress: ml::trainers::mamba2::TrainingProgress| {
if progress.epoch % 10 == 0 {
info!(
"📊 Epoch {}/{} ({:.1}%): loss={:.6}, perplexity={:.2}",
progress.epoch,
progress.total_epochs,
progress.progress_percentage,
progress.metrics.loss,
progress.metrics.perplexity
);
}
});
trainer.set_progress_callback(progress_callback);
// Train the model
info!("\n🏋️ Starting training...\n");
let start_time = std::time::Instant::now();
let training_history = trainer
.train(&train_data, &val_data)
.await
.context("Training failed")?;
let training_duration = start_time.elapsed();
// Print final metrics
info!("\n✅ Training completed successfully!");
info!("\n📊 Final Metrics:");
if let Some(final_epoch) = training_history.last() {
let loss_str = final_epoch
.loss
.map(|l| format!("{:.6}", l))
.unwrap_or_else(|| "N/A".to_string());
info!(" • Final loss: {}", loss_str);
let perplexity = final_epoch.loss.unwrap_or(f64::NAN).exp();
info!(" • Perplexity: {:.2}", perplexity);
}
info!(" • Best validation loss: {:.6}", trainer.best_val_loss);
info!(" • Epochs trained: {}", training_history.len());
info!(
" • Training time: {:.1}s ({:.1} min)",
training_duration.as_secs_f64(),
training_duration.as_secs_f64() / 60.0
);
// Get training statistics
let stats = trainer.get_training_statistics();
info!("\n📈 Training Statistics:");
if let Some(&memory_mb) = stats.get("estimated_memory_mb") {
info!(" • Memory usage: {:.1}MB", memory_mb);
}
if let Some(&throughput) = stats.get("throughput_pps") {
info!(" • Throughput: {:.0} predictions/sec", throughput);
}
info!(
"\n💾 Model checkpoints saved to: {}",
trainer.checkpoint_path
);
info!("\n🎉 MAMBA-2 training complete!");
Ok(())
}

406
crates/ml/examples/train_ppo.rs
View File

@@ -1,406 +0,0 @@
//! PPO Training Example with Real DataBento Market Data
//!
//! Trains a PPO model on real market data from DBN files with:
//! - Real OHLCV data + technical indicators
//! - Actual PnL-based rewards
//! - GAE advantages on real price trajectories
//! - Policy convergence validation (KL divergence > 0)
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (20 epochs)
//! cargo run -p ml --example train_ppo --release --features cuda
//!
//! # Custom epochs and output path
//! cargo run -p ml --example train_ppo --release --features cuda -- \
//! --epochs 50 \
//! --output-dir ml/trained_models \
//! --data-dir test_data/real/databento
//! ```
// Use mimalloc allocator for 10-25% performance improvement
#[cfg(feature = "mimalloc-allocator")]
use mimalloc::MiMalloc;
#[cfg(feature = "mimalloc-allocator")]
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;
use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;
use ml::data_loaders::BarSamplingMethod;
use ml::features::extraction::{extract_ml_features, OHLCVBar};
use ml::real_data_loader::RealDataLoader;
use ml::trainers::ppo::{PpoHyperparameters, PpoTrainer, PpoTrainingMetrics};
#[derive(Debug, Parser)]
#[command(name = "train_ppo", about = "Train PPO model on real market data")]
struct Opts {
/// Number of training epochs (default: 20 for policy convergence)
#[arg(long, default_value = "20")]
epochs: usize,
/// Learning rate
#[arg(long, default_value = "0.0003")]
learning_rate: f64,
/// Batch size (512 recommended for value network stability, prevents -23.56 explained variance failure)
#[arg(long, default_value = "512")]
batch_size: usize,
/// Output directory for trained model
#[arg(long, default_value = "ml/trained_models")]
output_dir: String,
/// Data directory containing DBN files
#[arg(long, default_value = "test_data/real/databento")]
data_dir: String,
/// Symbol to train on (ZN.FUT has ~29K bars)
#[arg(long, default_value = "ZN.FUT")]
symbol: String,
/// Verbose logging
#[arg(short, long)]
verbose: bool,
/// Enable early stopping (recommended, use --no-early-stopping to disable)
#[arg(long)]
early_stopping: bool,
/// Disable early stopping
#[arg(long)]
no_early_stopping: bool,
/// Minimum value loss improvement percentage for plateau detection
#[arg(long, default_value = "2.0")]
min_value_loss_improvement: f64,
/// Minimum explained variance threshold
#[arg(long, default_value = "0.4")]
min_explained_variance: f64,
/// Plateau detection window size (epochs)
#[arg(long, default_value = "30")]
plateau_window: usize,
/// Alternative bar sampling method (time, tick, volume, dollar, imbalance, run)
#[arg(long, default_value = "time")]
bar_method: String,
/// Bar sampling threshold (tick count, volume, dollar value, imbalance, or run length)
#[arg(long)]
bar_threshold: Option<f64>,
}
#[tokio::main]
async fn main() -> Result<()> {
// Parse CLI options
let opts = Opts::parse();
// Setup logging
let level = if opts.verbose {
tracing::Level::DEBUG
} else {
tracing::Level::INFO
};
let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
tracing::subscriber::set_global_default(subscriber)
.context("Failed to set tracing subscriber")?;
#[cfg(feature = "mimalloc-allocator")]
info!("🚀 Using mimalloc allocator for improved performance");
#[cfg(not(feature = "mimalloc-allocator"))]
info!("ℹ️ Using system allocator (consider --features mimalloc-allocator for 10-25% speedup)");
info!("🚀 Starting PPO Training with Real DataBento Data");
info!("Configuration:");
info!(" • Epochs: {}", opts.epochs);
info!(" • Learning rate: {}", opts.learning_rate);
info!(" • Batch size: {}", opts.batch_size);
info!(" • GPU: CUDA MANDATORY (no CPU fallback)");
info!(" • Output directory: {}", opts.output_dir);
info!(" • Data directory: {}", opts.data_dir);
info!(" • Symbol: {}", opts.symbol);
info!(" • Bar sampling method: {}", opts.bar_method);
if let Some(threshold) = opts.bar_threshold {
info!(" • Bar threshold: {}", threshold);
}
// Determine early stopping (enabled by default, unless --no-early-stopping is specified)
let early_stopping_enabled = !opts.no_early_stopping;
info!(
" • Early stopping: {}",
if early_stopping_enabled {
"enabled"
} else {
"disabled"
}
);
if early_stopping_enabled {
info!(
" - Min value loss improvement: {}%",
opts.min_value_loss_improvement
);
info!(
" - Min explained variance: {}",
opts.min_explained_variance
);
info!(" - Plateau window: {} epochs", opts.plateau_window);
}
// Create output directory
let output_path = PathBuf::from(&opts.output_dir);
if !output_path.exists() {
std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
info!("✅ Created output directory: {}", opts.output_dir);
}
// Configure alternative bar sampling (Wave B)
let bar_sampling = match opts.bar_method.as_str() {
"tick" => BarSamplingMethod::TickBars(opts.bar_threshold.unwrap_or(100.0) as usize),
"volume" => BarSamplingMethod::VolumeBars(opts.bar_threshold.unwrap_or(10000.0)),
"dollar" => BarSamplingMethod::DollarBars(opts.bar_threshold.unwrap_or(2_000_000.0)),
"imbalance" => BarSamplingMethod::ImbalanceBars(opts.bar_threshold.unwrap_or(1000.0)),
"run" => BarSamplingMethod::RunBars(opts.bar_threshold.unwrap_or(50.0) as usize),
_ => BarSamplingMethod::TimeBars,
};
info!("✅ Bar sampling configured: {:?}", bar_sampling);
// Load real market data from DBN files
info!("\n📊 Loading real market data from DBN files...");
let mut loader = RealDataLoader::new(&opts.data_dir);
// Note: RealDataLoader will need to accept bar_sampling parameter
// This requires updating RealDataLoader to use alternative bar sampling
let bars = loader
.load_symbol_data(&opts.symbol)
.await
.context(format!("Failed to load data for symbol: {}", opts.symbol))?;
info!("✅ Loaded {} OHLCV bars for {}", bars.len(), opts.symbol);
// Extract features and indicators
info!("\n🔧 Extracting features and technical indicators...");
let features = loader
.extract_features(&bars)
.context("Failed to extract features")?;
let _indicators = loader
.calculate_indicators(&bars)
.context("Failed to calculate indicators")?;
info!("✅ Feature extraction complete:");
info!(" • OHLCV bars: {}", features.prices.len());
info!(" • Returns: {}", features.returns.len());
info!(" • Volume: {}", features.volume.len());
info!(" • Indicators: 10 technical indicators");
// Build PPO state vectors using 54-feature extraction pipeline
// Features 0-4: OHLCV (normalized)
// Features 5-14: Technical indicators (10)
// Remaining features from extraction pipeline
info!("\n🏗️ Extracting 54-dimensional feature vectors...");
// Convert RealDataLoader bars to OHLCVBar format for feature extraction
let ohlcv_bars: Vec<OHLCVBar> = bars
.iter()
.map(|bar| OHLCVBar {
timestamp: bar.timestamp,
open: bar.open,
high: bar.high,
low: bar.low,
close: bar.close,
volume: bar.volume,
})
.collect();
// Extract 54-dimensional feature vectors (requires 50-bar warmup)
let feature_vectors =
extract_ml_features(&ohlcv_bars).context("Failed to extract 54-dimensional features")?;
info!(
"✅ Extracted {} feature vectors (dim=54, warmup bars skipped=50)",
feature_vectors.len()
);
// Convert FeatureVector ([f64; 54]) to Vec<Vec<f32>> for PPO trainer
let state_dim = 54;
let market_data: Vec<Vec<f32>> = feature_vectors
.iter()
.map(|fv| fv.iter().map(|&v| v as f32).collect())
.collect();
// Validate state dimensions
if let Some(first_state) = market_data.first() {
if first_state.len() != state_dim {
return Err(anyhow::anyhow!(
"State dimension mismatch: expected {}, got {}",
state_dim,
first_state.len()
));
}
}
// Configure PPO hyperparameters
let hyperparams = PpoHyperparameters {
learning_rate: opts.learning_rate,
batch_size: opts.batch_size,
gamma: 0.99,
clip_epsilon: 0.2,
vf_coef: 0.5,
ent_coef: 0.01,
gae_lambda: 0.95,
rollout_steps: 2048,
minibatch_size: opts.batch_size,
epochs: opts.epochs,
early_stopping_enabled,
min_value_loss_improvement_pct: opts.min_value_loss_improvement,
min_explained_variance: opts.min_explained_variance,
plateau_window: opts.plateau_window,
min_epochs_before_stopping: 50,
};
// Create PPO trainer with real data state dimension
let trainer = PpoTrainer::new(
hyperparams.clone(),
state_dim,
&opts.output_dir,
true, // CUDA always required
None, // Single environment (standard mode)
)
.context("Failed to create PPO trainer")?;
info!("✅ PPO trainer initialized (state_dim={})", state_dim);
// Create progress callback with convergence tracking
let mut policy_updates = 0;
let mut kl_divergence_history = Vec::new();
let progress_callback = |metrics: PpoTrainingMetrics| {
// Track policy updates (KL divergence > 0 indicates policy changed)
if metrics.kl_divergence > 0.0 {
policy_updates += 1;
}
kl_divergence_history.push(metrics.kl_divergence);
info!(
"📊 Epoch {}/{}: policy_loss={:.4}, value_loss={:.4}, kl_div={:.6}, expl_var={:.4}, mean_reward={:.4}",
metrics.epoch,
hyperparams.epochs,
metrics.policy_loss,
metrics.value_loss,
metrics.kl_divergence,
metrics.explained_variance,
metrics.mean_reward
);
};
// Train the model
info!("\n🏋️ Starting training...\n");
let start_time = std::time::Instant::now();
let final_metrics = trainer
.train(market_data, progress_callback)
.await
.context("Training failed")?;
let training_duration = start_time.elapsed();
// Print final metrics
info!("\n✅ Training completed successfully!");
info!("\n📊 Final Metrics:");
info!(" • Policy loss: {:.6}", final_metrics.policy_loss);
info!(" • Value loss: {:.6}", final_metrics.value_loss);
info!(" • KL divergence: {:.6}", final_metrics.kl_divergence);
info!(
" • Explained variance: {:.4}",
final_metrics.explained_variance
);
info!(" • Mean reward: {:.4}", final_metrics.mean_reward);
info!(" • Std reward: {:.4}", final_metrics.std_reward);
info!(" • Entropy: {:.4}", final_metrics.entropy);
info!(
" • Training time: {:.1}s ({:.1} min)",
training_duration.as_secs_f64(),
training_duration.as_secs_f64() / 60.0
);
// Validate policy convergence
info!("\n🔍 Policy Convergence Analysis:");
info!(" • Total epochs: {}", hyperparams.epochs);
info!(" • Policy updates (KL > 0): {}", policy_updates);
info!(
" • Policy update rate: {:.1}%",
(policy_updates as f64 / hyperparams.epochs as f64) * 100.0
);
// Calculate KL divergence statistics
let kl_mean = kl_divergence_history.iter().sum::<f32>() / kl_divergence_history.len() as f32;
let kl_max = kl_divergence_history
.iter()
.copied()
.fold(f32::NEG_INFINITY, f32::max);
let kl_min = kl_divergence_history
.iter()
.copied()
.fold(f32::INFINITY, f32::min);
info!(" • KL divergence (mean): {:.6}", kl_mean);
info!(" • KL divergence (max): {:.6}", kl_max);
info!(" • KL divergence (min): {:.6}", kl_min);
// Convergence validation
if final_metrics.kl_divergence > 0.0 {
info!(" ✅ PASS: Policy updates detected (KL divergence > 0)");
} else {
warn!(" ⚠️ WARN: No policy updates in final epoch (KL divergence = 0)");
warn!(" This may indicate learning rate too low or convergence");
}
// Value function validation
if final_metrics.explained_variance > 0.5 {
info!(" ✅ PASS: Value network learning (explained variance > 0.5)");
} else {
warn!(" ⚠️ WARN: Value network may need tuning (explained variance < 0.5)");
}
// Checkpoint is already saved by trainer (every 10 epochs)
let final_checkpoint = output_path.join(format!(
"ppo_checkpoint_epoch_{}.safetensors",
hyperparams.epochs
));
info!(
"\n💾 Final checkpoint saved to: {}",
final_checkpoint.display()
);
info!("\n🎉 PPO training complete with real DataBento data!");
info!("📁 Model files saved to: {}", opts.output_dir);
info!("\n📈 Training Summary:");
info!(" • Data source: Real DataBento OHLCV ({})", opts.symbol);
info!(" • Training samples: {}", bars.len());
info!(" • State dimension: {}", state_dim);
info!(" • Features: OHLCV + 10 technical indicators + log returns");
info!(
" • Policy updates: {}/{} epochs ({:.1}%)",
policy_updates,
hyperparams.epochs,
(policy_updates as f64 / hyperparams.epochs as f64) * 100.0
);
info!(
" • Convergence: {}",
if final_metrics.kl_divergence > 0.0 {
"✅ Achieved"
} else {
"⚠️ Check logs"
}
);
Ok(())
}

16
scripts/entrypoint-generic.sh
View File

@@ -52,17 +52,17 @@ if [ $# -eq 0 ]; then
log " Override Docker CMD with your training command:"
log ""
log "EXAMPLES:"
log " # DQN Training (100 epochs)"
log " CMD: /runpod-volume/binaries/train_dqn --parquet-file /runpod-volume/test_data/ES_FUT_small.parquet --epochs 100 --output-dir /runpod-volume/models"
log " # RL Training (DQN on ES futures, walk-forward)"
log " CMD: /runpod-volume/binaries/train_baseline_rl --model dqn --symbol ES.FUT --data-dir /runpod-volume/data --epochs 100"
log ""
log " # TFT Training (50 epochs, GPU enabled)"
log " CMD: /runpod-volume/binaries/train_tft_parquet --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --use-gpu --output-dir /runpod-volume/models"
log " # RL Training (PPO on NQ futures)"
log " CMD: /runpod-volume/binaries/train_baseline_rl --model ppo --symbol NQ.FUT --data-dir /runpod-volume/data --epochs 200"
log ""
log " # PPO Training (200 epochs)"
log " CMD: /runpod-volume/binaries/train_ppo --parquet-file /runpod-volume/test_data/NQ_FUT_180d.parquet --epochs 200 --output-dir /runpod-volume/models"
log " # Supervised Training (TFT on parquet data)"
log " CMD: /runpod-volume/binaries/train_baseline_supervised --model tft --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50"
log ""
log " # MAMBA-2 Training (30 epochs)"
log " CMD: /runpod-volume/binaries/train_mamba2_parquet --parquet-file /runpod-volume/test_data/6E_FUT_180d.parquet --epochs 30 --output-dir /runpod-volume/models"
log " # Supervised Training (Mamba2 on parquet data)"
log " CMD: /runpod-volume/binaries/train_baseline_supervised --model mamba2 --parquet-file /runpod-volume/test_data/6E_FUT_180d.parquet --epochs 30"
log ""
log "Container is ready and waiting for commands..."
log "To deploy with a custom command, override the Docker CMD in your deployment script"