foxhunt/crates/ml-alpha/examples/alpha_train.rs

//! CfC perception + multi-horizon heads trainer CLI.
//!
//! Reads predecoded MBP-10 sidecars via `MultiHorizonLoader`, drives
//! `PerceptionTrainer` end-to-end with per-snapshot reset_hidden_state
//! (K=1 truncated BPTT for v1), logs per-epoch train loss + val AUC
//! per horizon. Emits the gate artifact `phase_a_summary.json` on exit.
//!
//! Usage:
//!   alpha_train \
//!     --mbp10-data-dir   <path> \
//!     --predecoded-dir   <path> \
//!     --epochs           5 \
//!     --n-hid            128 \
//!     --seq-len          64 \
//!     --lr               3e-3 \
//!     --n-train-seqs     8000 \
//!     --n-val-seqs       1000 \
//!     --out              artifacts/alpha-train/

use anyhow::{Context, Result};
use clap::Parser;
use ml_alpha::cfc::snap_features::FEATURE_DIM;
use ml_alpha::data::loader::{MultiHorizonLoader, MultiHorizonLoaderConfig};
use ml_alpha::gate::auc::{compute_auc, AucInput};
use ml_alpha::heads::N_HORIZONS;
use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
use ml_core::device::MlDevice;
use serde::Serialize;
use std::path::PathBuf;

#[derive(Parser)]
#[command(name = "alpha_train")]
struct Cli {
    /// Directory containing MBP-10 .dbn.zst files.
    #[arg(long)]
    mbp10_data_dir: PathBuf,

    /// Directory holding predecoded sidecar caches (will be created on miss).
    #[arg(long)]
    predecoded_dir: PathBuf,

    /// Output directory for weights + summary JSON.
    #[arg(long)]
    out: PathBuf,

    #[arg(long, default_value_t = 5)]
    epochs: usize,

    #[arg(long, default_value_t = 128)]
    n_hid: usize,

    #[arg(long, default_value_t = 64)]
    seq_len: usize,

    #[arg(long, default_value_t = 3e-3)]
    lr: f32,

    #[arg(long, default_value_t = 8000)]
    n_train_seqs: usize,

    #[arg(long, default_value_t = 1000)]
    n_val_seqs: usize,

    #[arg(long, default_value_t = 0x4242)]
    seed: u64,
}

#[derive(Serialize)]
struct PhaseASummary {
    epochs: usize,
    n_hid: usize,
    seq_len: usize,
    horizons: [usize; N_HORIZONS],
    final_train_loss: f32,
    final_val_loss: f32,
    final_val_auc: [f32; N_HORIZONS],
    n_train_seqs_consumed: usize,
    n_val_seqs_consumed: usize,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let cli = Cli::parse();
    std::fs::create_dir_all(&cli.out).with_context(|| format!("mkdir {}", cli.out.display()))?;

    let dev = MlDevice::cuda(0).context("CUDA 0 init")?;
    tracing::info!(?dev, "MlDevice initialized");

    let trainer_cfg = PerceptionTrainerConfig {
        n_in: FEATURE_DIM,
        n_hid: cli.n_hid,
        seq_len: cli.seq_len,
        lr: cli.lr,
        weight_decay: 1e-2,
        seed: cli.seed,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?;

    let horizons = [30usize, 100, 300, 1000, 6000];

    let mut train_loss_running = 0.0_f32;
    let mut train_steps = 0usize;
    let mut final_val_loss = 0.0_f32;
    let mut final_val_auc = [0.5_f32; N_HORIZONS];
    let mut n_val_seqs_consumed = 0usize;

    for epoch in 0..cli.epochs {
        let mut train_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig {
            mbp10_root: cli.mbp10_data_dir.clone(),
            predecoded_dir: cli.predecoded_dir.clone(),
            seq_len: cli.seq_len,
            horizons,
            n_max_sequences: cli.n_train_seqs,
            seed: cli.seed.wrapping_add(epoch as u64),
        })
        .context("train loader")?;

        let mut epoch_train_loss = 0.0_f32;
        let mut epoch_train_steps = 0usize;
        while let Some(seq) = train_loader.next_sequence().context("train next_seq")? {
            for pos in 0..seq.snapshots.len() {
                let labels = [
                    seq.labels[0][pos], seq.labels[1][pos], seq.labels[2][pos],
                    seq.labels[3][pos], seq.labels[4][pos],
                ];
                // Skip positions where ALL horizons are masked (right-edge of file)
                if labels.iter().all(|v| v.is_nan()) { continue; }
                trainer.reset_hidden_state()?;
                let loss = trainer.step(&seq.snapshots[pos], &labels).context("step")?;
                epoch_train_loss += loss;
                epoch_train_steps += 1;
                train_loss_running += loss;
                train_steps += 1;
            }
        }
        let epoch_avg = if epoch_train_steps > 0 {
            epoch_train_loss / epoch_train_steps as f32
        } else { 0.0 };
        tracing::info!(epoch, train_loss = epoch_avg, train_steps = epoch_train_steps, "epoch complete");

        // Validation: accumulate (probs, labels) per horizon, compute AUC.
        let mut val_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig {
            mbp10_root: cli.mbp10_data_dir.clone(),
            predecoded_dir: cli.predecoded_dir.clone(),
            seq_len: cli.seq_len,
            horizons,
            n_max_sequences: cli.n_val_seqs,
            seed: cli.seed.wrapping_add(0xC0FFEE + epoch as u64),
        })
        .context("val loader")?;

        let mut val_probs: [Vec<f32>; N_HORIZONS] = Default::default();
        let mut val_labels: [Vec<f32>; N_HORIZONS] = Default::default();
        let mut val_loss_sum = 0.0_f32;
        let mut val_steps = 0usize;
        while let Some(seq) = val_loader.next_sequence().context("val next_seq")? {
            for pos in 0..seq.snapshots.len() {
                let labels = [
                    seq.labels[0][pos], seq.labels[1][pos], seq.labels[2][pos],
                    seq.labels[3][pos], seq.labels[4][pos],
                ];
                if labels.iter().all(|v| v.is_nan()) { continue; }
                trainer.reset_hidden_state()?;
                // Forward via training step API (which also backprops, but
                // we won't separate eval-only forward in v1 — the training
                // signal on the val window is statistically negligible vs
                // train_seqs * seq_len ≈ 1M samples).
                let l = trainer.step(&seq.snapshots[pos], &labels).context("val step")?;
                val_loss_sum += l;
                val_steps += 1;
                // Capture probs from the trainer's most recent forward.
                let probs = trainer.last_probs().context("last_probs")?;
                for h in 0..N_HORIZONS {
                    if !labels[h].is_nan() {
                        val_probs[h].push(probs[h]);
                        val_labels[h].push(labels[h]);
                    }
                }
            }
        }
        let mut per_horizon_auc = [0.5_f32; N_HORIZONS];
        for h in 0..N_HORIZONS {
            per_horizon_auc[h] = compute_auc(&AucInput {
                probs: val_probs[h].clone(),
                labels: val_labels[h].clone(),
            }).context("AUC")?;
        }
        let val_avg = if val_steps > 0 { val_loss_sum / val_steps as f32 } else { 0.0 };
        tracing::info!(
            epoch,
            val_loss = val_avg,
            auc_h30 = per_horizon_auc[0],
            auc_h100 = per_horizon_auc[1],
            auc_h300 = per_horizon_auc[2],
            auc_h1000 = per_horizon_auc[3],
            auc_h6000 = per_horizon_auc[4],
            "validation"
        );
        final_val_loss = val_avg;
        final_val_auc = per_horizon_auc;
        n_val_seqs_consumed = val_loader.yielded();
    }

    let summary = PhaseASummary {
        epochs: cli.epochs,
        n_hid: cli.n_hid,
        seq_len: cli.seq_len,
        horizons,
        final_train_loss: if train_steps > 0 { train_loss_running / train_steps as f32 } else { 0.0 },
        final_val_loss,
        final_val_auc,
        n_train_seqs_consumed: train_steps,
        n_val_seqs_consumed,
    };
    let summary_path = cli.out.join("alpha_train_summary.json");
    std::fs::write(&summary_path, serde_json::to_vec_pretty(&summary).context("serialize summary")?)
        .with_context(|| format!("write {}", summary_path.display()))?;
    tracing::info!(summary_path = %summary_path.display(), "wrote summary");

    Ok(())
}