foxhunt/crates/ml-dqn/src/dueling.rs

//! Dueling Q-Network Architecture
//!
//! Implements the Dueling DQN architecture from "Dueling Network Architectures for Deep
//! Reinforcement Learning" (Wang et al., 2016).
//!
//! ## Architecture
//!
//! ```text
//! State [state_dim] -> Shared Features [hidden_dim]
//!                      |                |
//!                   Value V(s)      Advantage A(s,a)
//!                   [1 scalar]      [num_actions]
//!                      |                |
//!                   Q(s,a) = V(s) + A(s,a) - mean(A(s,.))
//! ```
//!
//! ## Key Features
//!
//! - **Separate Value/Advantage Streams**: Decomposes Q-values into state value and advantage
//! - **Mean Subtraction**: Ensures identifiability (zero mean advantage)
//! - **Improved Learning**: Better gradient flow and sample efficiency (+10-20%)
//! - **Numerical Stability**: Xavier initialization for all layers
//!
//! ## Mathematical Formulation
//!
//! `Q(s,a) = V(s) + [A(s,a) - (1/|A|) * sum_a' A(s,a')]`
//!
//! Where:
//! - V(s): State value function (scalar)
//! - A(s,a): Advantage function (per-action)
//! - mean(A(s,.)): Average advantage across all actions (ensures zero mean)

use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;
use serde::{Deserialize, Serialize};

use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore};
use ml_core::MLError;

/// Configuration for Dueling Q-Network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DuelingConfig {
    /// Input state dimension
    pub state_dim: usize,
    /// Number of actions
    pub num_actions: usize,
    /// Shared feature layer dimensions
    pub shared_hidden_dims: Vec<usize>,
    /// Value stream hidden dimension
    pub value_hidden_dim: usize,
    /// Advantage stream hidden dimension
    pub advantage_hidden_dim: usize,
    /// `LeakyReLU` negative slope
    pub leaky_relu_alpha: f64,
    /// Dropout rate applied after shared layer activations (0.0 = no dropout)
    #[serde(default)]
    pub dropout_rate: f64,
}

impl DuelingConfig {
    /// Create default dueling config from basic parameters
    pub const fn new(
        state_dim: usize,
        num_actions: usize,
        shared_hidden_dims: Vec<usize>,
        value_hidden_dim: usize,
        advantage_hidden_dim: usize,
    ) -> Self {
        Self {
            state_dim,
            num_actions,
            shared_hidden_dims,
            value_hidden_dim,
            advantage_hidden_dim,
            leaky_relu_alpha: 0.01,
            dropout_rate: 0.0,
        }
    }

    /// Create from `DQNConfig` parameters
    ///
    /// All `hidden_dims` become shared feature layers. The value and advantage
    /// heads branch off after the last shared layer, each with its own
    /// `dueling_hidden_dim`-wide hidden layer. This matches the CUDA
    /// experience-collection kernel which hardcodes 2 shared layers.
    pub fn from_dqn_params(
        state_dim: usize,
        num_actions: usize,
        hidden_dims: &[usize],
        dueling_hidden_dim: usize,
        leaky_relu_alpha: f64,
    ) -> Self {
        Self {
            state_dim,
            num_actions,
            shared_hidden_dims: hidden_dims.to_vec(),
            value_hidden_dim: dueling_hidden_dim,
            advantage_hidden_dim: dueling_hidden_dim,
            leaky_relu_alpha,
            dropout_rate: 0.0,
        }
    }
}

/// Dueling Q-Network with separate value and advantage streams.
///
/// Weights stored in `GpuVarStore` with `GpuLinear` layers using cuBLAS sgemm.
/// Forward pass runs entirely on GPU (cold path downloads to host for activation/mean ops).
#[allow(missing_debug_implementations)]
pub struct DuelingQNetwork {
    /// Shared feature extraction layers
    shared_layers: Vec<GpuLinear>,

    /// Value stream layers
    value_fc: GpuLinear,
    value_out: GpuLinear, // Output: [batch, 1]

    /// Advantage stream layers
    advantage_fc: GpuLinear,
    advantage_out: GpuLinear, // Output: [batch, num_actions]

    /// Configuration
    config: DuelingConfig,

    /// Native CUDA weight storage
    store: GpuVarStore,

    /// cuBLAS handle for sgemm
    cublas: CudaBlas,

    /// CUDA stream
    stream: Arc<CudaStream>,
}

impl DuelingQNetwork {
    /// Create new dueling Q-network
    ///
    /// # Arguments
    ///
    /// * `config` - Dueling network configuration
    /// * `stream` - CUDA stream for all operations
    ///
    /// # Returns
    ///
    /// New `DuelingQNetwork` instance with Xavier-initialized weights
    pub fn new(config: DuelingConfig, stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let mut store = GpuVarStore::new(Arc::clone(&stream));

        // Build shared feature layers
        let mut shared_layers = Vec::new();
        let mut current_dim = config.state_dim;

        for (i, &hidden_dim) in config.shared_hidden_dims.iter().enumerate() {
            let layer = store.linear(
                &format!("shared_{i}"),
                current_dim,
                hidden_dim,
            )?;
            shared_layers.push(layer);
            current_dim = hidden_dim;
        }

        // Value stream
        let value_fc = store.linear("value_fc", current_dim, config.value_hidden_dim)?;
        let value_out = store.linear("value_out", config.value_hidden_dim, 1)?;

        // Advantage stream
        let advantage_fc = store.linear("advantage_fc", current_dim, config.advantage_hidden_dim)?;
        let advantage_out = store.linear("advantage_out", config.advantage_hidden_dim, config.num_actions)?;

        let cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| MLError::ModelError(format!("cuBLAS init: {e}")))?;

        Ok(Self {
            shared_layers,
            value_fc,
            value_out,
            advantage_fc,
            advantage_out,
            config,
            store,
            cublas,
            stream,
        })
    }

    /// Forward pass through dueling network (cold path).
    ///
    /// # Arguments
    ///
    /// * `state` - State data as flat f32 slice, shape [`batch_size` * `state_dim`]
    /// * `batch_size` - Number of samples in the batch
    ///
    /// # Returns
    ///
    /// Q-values as flat f32 Vec, shape [`batch_size` * `num_actions`]
    ///
    /// # Mathematical Formula
    ///
    /// Q(s,a) = V(s) + [A(s,a) - mean(A(s,.))]
    pub fn forward(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        let state_dim = self.config.state_dim;
        let num_actions = self.config.num_actions;
        let alpha = self.config.leaky_relu_alpha as f32;

        // Upload state to GPU
        let mut h = GpuTensor::from_host(state, vec![batch_size, state_dim], &self.stream)?;

        // Shared feature extraction with LeakyReLU
        for layer in &self.shared_layers {
            let (out, _acts) = layer.forward(&h, &self.store, &self.cublas, &self.stream)?;
            // LeakyReLU on host (cold path)
            let mut host = out.to_host(&self.stream)?;
            for v in host.iter_mut() {
                if *v < 0.0 { *v *= alpha; }
            }
            h = GpuTensor::from_host(&host, out.shape().to_vec(), &self.stream)?;
        }

        // Value stream: V(s) -> [batch, 1]
        let (v_hidden, _) = self.value_fc.forward(&h, &self.store, &self.cublas, &self.stream)?;
        let mut v_host = v_hidden.to_host(&self.stream)?;
        for v in v_host.iter_mut() { if *v < 0.0 { *v *= alpha; } }
        let v_gpu = GpuTensor::from_host(&v_host, vec![batch_size, self.config.value_hidden_dim], &self.stream)?;
        let (v_out_gpu, _) = self.value_out.forward(&v_gpu, &self.store, &self.cublas, &self.stream)?;
        let v_out = v_out_gpu.to_host(&self.stream)?; // [batch_size * 1]

        // Advantage stream: A(s,a) -> [batch, num_actions]
        let (a_hidden, _) = self.advantage_fc.forward(&h, &self.store, &self.cublas, &self.stream)?;
        let mut a_host = a_hidden.to_host(&self.stream)?;
        for v in a_host.iter_mut() { if *v < 0.0 { *v *= alpha; } }
        let a_gpu = GpuTensor::from_host(&a_host, vec![batch_size, self.config.advantage_hidden_dim], &self.stream)?;
        let (a_out_gpu, _) = self.advantage_out.forward(&a_gpu, &self.store, &self.cublas, &self.stream)?;
        let a_out = a_out_gpu.to_host(&self.stream)?; // [batch_size * num_actions]

        // Combine: Q(s,a) = V(s) + A(s,a) - mean(A(s,.))
        let mut q_values = vec![0.0_f32; batch_size * num_actions];
        for b in 0..batch_size {
            let v = v_out[b]; // V(s) scalar for this sample

            // Compute mean advantage for this sample
            let a_start = b * num_actions;
            let a_slice = &a_out[a_start..a_start + num_actions];
            let a_mean: f32 = a_slice.iter().sum::<f32>() / num_actions as f32;

            // Q(s,a) = V(s) + A(s,a) - mean(A)
            for a in 0..num_actions {
                q_values[b * num_actions + a] = v + a_slice[a] - a_mean;
            }
        }

        Ok(q_values)
    }

    /// Forward pass for a single state (convenience).
    pub fn forward_single(&self, state: &[f32]) -> Result<Vec<f32>, MLError> {
        self.forward(state, 1)
    }

    /// Get `GpuVarStore` for weight serialization
    pub const fn store(&self) -> &GpuVarStore {
        &self.store
    }

    /// Get mutable `GpuVarStore` for weight updates
    pub const fn store_mut(&mut self) -> &mut GpuVarStore {
        &mut self.store
    }

    /// Get configuration
    pub const fn config(&self) -> &DuelingConfig {
        &self.config
    }

    /// Copy weights from another dueling network
    pub fn copy_weights_from(&mut self, other: &DuelingQNetwork) -> Result<(), MLError> {
        let exported = other.store.export_to_host()?; // gpu-exit: checkpoint weight export
        self.store.import_from_host(&exported)
    }
}

#[allow(clippy::items_after_test_module)]
#[cfg(test)]
#[allow(clippy::unnecessary_wraps)]
mod tests {
    use super::*;

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
            .expect("CUDA required")
            .new_stream()
            .expect("CUDA stream")
    }

    #[test]
    fn test_dueling_network_creation() -> anyhow::Result<()> {
        let config = DuelingConfig::new(
            32,              // state_dim
            5,               // num_actions (5 exposure levels)
            vec![256, 128],  // shared_hidden_dims
            64,              // value_hidden_dim
            64,              // advantage_hidden_dim
        );

        let network = DuelingQNetwork::new(config, make_stream())?;

        assert_eq!(network.shared_layers.len(), 2);
        Ok(())
    }

    #[test]
    fn test_dueling_forward_pass() -> anyhow::Result<()> {
        let config = DuelingConfig::new(32, 5, vec![256, 128], 64, 64);
        let stream = make_stream();
        let network = DuelingQNetwork::new(config, Arc::clone(&stream))?;

        // Create batch of states
        let batch_size = 4;
        let state_data: Vec<f32> = (0..batch_size * 32).map(|i| (i as f32 * 0.01).sin()).collect();

        // Forward pass
        let q_values = network.forward(&state_data, batch_size)?;

        // Check output length
        assert_eq!(q_values.len(), batch_size * 5);

        Ok(())
    }

    #[test]
    fn test_dueling_mean_subtraction() -> anyhow::Result<()> {
        // Test that mean(A) is correctly subtracted, ensuring zero-mean advantage
        let config = DuelingConfig::new(4, 3, vec![8], 4, 4);
        let stream = make_stream();
        let network = DuelingQNetwork::new(config, Arc::clone(&stream))?;

        // Simple state
        let state = vec![1.0_f32; 4];

        // Forward pass
        let q_values = network.forward_single(&state)?;

        // Q-values should be valid (no NaN/Inf)
        for &q in &q_values {
            assert!(q.is_finite(), "Q-value should be finite, got {}", q);
        }

        Ok(())
    }

    #[test]
    fn test_dueling_from_dqn_params() -> anyhow::Result<()> {
        let config = DuelingConfig::from_dqn_params(
            32,              // state_dim
            5,               // num_actions (5 exposure levels)
            &[256, 128, 64], // hidden_dims
            64,              // dueling_hidden_dim
            0.01,            // leaky_relu_alpha
        );

        assert_eq!(config.state_dim, 32);
        assert_eq!(config.num_actions, 5);
        assert_eq!(config.shared_hidden_dims, vec![256, 128, 64]); // All hidden dims become shared layers
        assert_eq!(config.value_hidden_dim, 64);
        assert_eq!(config.advantage_hidden_dim, 64);

        Ok(())
    }

    #[test]
    fn test_dueling_weight_copy() -> anyhow::Result<()> {
        let config = DuelingConfig::new(8, 3, vec![16], 8, 8);
        let stream = make_stream();

        let network1 = DuelingQNetwork::new(config.clone(), Arc::clone(&stream))?;
        let mut network2 = DuelingQNetwork::new(config, Arc::clone(&stream))?;

        // Copy weights
        network2.copy_weights_from(&network1)?;

        // Verify same output for same input
        let state = vec![1.0_f32; 8];
        let q1 = network1.forward_single(&state)?;
        let q2 = network2.forward_single(&state)?;

        for (v1, v2) in q1.iter().zip(q2.iter()) {
            assert!((v1 - v2).abs() < 1e-5, "Q-values should match after copy");
        }

        Ok(())
    }
}

// Manual Debug implementation for DuelingQNetwork
impl std::fmt::Debug for DuelingQNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DuelingQNetwork")
            .field("config", &self.config)
            .field("num_shared_layers", &self.shared_layers.len())
            .finish()
    }
}