foxhunt/ml/examples/measure_dqn_memory.rs

//! Simple DQN memory measurement tool
//!
//! This script measures the GPU memory footprint of the DQN model
//! and provides baseline metrics for optimization.

use candle_core::Device;
use ml::dqn::{WorkingDQN, WorkingDQNConfig};

fn main() -> anyhow::Result<()> {
    // Initialize device
    let device = Device::cuda_if_available(0)?;
    println!(
        "Device: {:?}",
        if device.is_cuda() { "CUDA" } else { "CPU" }
    );

    // Measure baseline memory
    #[cfg(feature = "cuda")]
    {
        use std::process::Command;

        let output = Command::new("nvidia-smi")
            .args(&["--query-gpu=memory.used", "--format=csv,noheader,nounits"])
            .output()?;

        let baseline_mb: f64 = String::from_utf8_lossy(&output.stdout)
            .trim()
            .parse()
            .unwrap_or(0.0);

        println!("Baseline GPU memory: {:.0} MB", baseline_mb);

        // Create DQN config (54 features)
        let config = WorkingDQNConfig {
            state_dim: 54,
            num_actions: 3,
            hidden_dims: vec![128, 64, 32],
            learning_rate: 0.0001,
            gamma: 0.99,
            epsilon_start: 1.0,
            epsilon_end: 0.01,
            epsilon_decay: 0.995,
            replay_buffer_capacity: 100_000,
            batch_size: 128,
            min_replay_size: 256,
            target_update_freq: 1000,
            use_double_dqn: true,
            use_huber_loss: true, // Huber loss default (more robust to outliers)
            huber_delta: 1.0,     // Standard Huber delta
        };

        println!("\nCreating DQN model...");
        let dqn = WorkingDQN::new(config)?;

        // Measure after model creation
        let output = Command::new("nvidia-smi")
            .args(&["--query-gpu=memory.used", "--format=csv,noheader,nounits"])
            .output()?;

        let model_mb: f64 = String::from_utf8_lossy(&output.stdout)
            .trim()
            .parse()
            .unwrap_or(0.0);

        let dqn_memory = model_mb - baseline_mb;

        println!("\n=== DQN MEMORY REPORT ===");
        println!("DQN Model Memory: {:.0} MB", dqn_memory);
        println!("Target: <150 MB");
        println!(
            "Status: {}",
            if dqn_memory <= 150.0 {
                "✅ PASS"
            } else {
                "❌ FAIL"
            }
        );
        println!();

        // Model details
        println!("Model Configuration:");
        println!("  State dimension: 54");
        println!("  Hidden layers: [128, 64, 32]");
        println!("  Output actions: 3");
        println!("  Replay buffer: 100,000");
        println!("  Double DQN: enabled");

        // Calculate theoretical parameter count
        let params = (54 * 128) + 128 +  // input -> hidden1
                     (128 * 64) + 64 +     // hidden1 -> hidden2
                     (64 * 32) + 32 +      // hidden2 -> hidden3
                     (32 * 3) + 3; // hidden3 -> output

        let params_mb = (params * 4) as f64 / 1024.0 / 1024.0; // FP32

        println!("\nTheoretical Model Size:");
        println!("  Parameters: {}", params);
        println!("  FP32 size: {:.2} MB", params_mb);
        println!("  Actual GPU memory: {:.0} MB", dqn_memory);
        println!(
            "  Overhead: {:.0} MB ({:.1}%)",
            dqn_memory - params_mb,
            ((dqn_memory - params_mb) / dqn_memory) * 100.0
        );

        // Don't drop DQN to avoid deallocation before measurement
        std::mem::forget(dqn);
    }

    #[cfg(not(feature = "cuda"))]
    {
        println!("CUDA not available - memory measurement requires GPU");
    }

    Ok(())
}