foxhunt/ml/src/dqn/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|) * Σ_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 candle_core::{DType, Device, Tensor};
use candle_nn::{Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};

use crate::dqn::xavier_init::linear_xavier;
use crate::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,
}

impl DuelingConfig {
    /// Create default dueling config from basic parameters
    pub 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,
        }
    }

    /// Create from DQNConfig parameters
    pub fn from_dqn_params(
        state_dim: usize,
        num_actions: usize,
        hidden_dims: &[usize],
        dueling_hidden_dim: usize,
        leaky_relu_alpha: f64,
    ) -> Self {
        // Use first N-1 layers as shared features, split at final layer
        let shared_dims = if hidden_dims.len() > 1 {
            hidden_dims[..hidden_dims.len() - 1].to_vec()
        } else {
            hidden_dims.to_vec()
        };

        Self {
            state_dim,
            num_actions,
            shared_hidden_dims: shared_dims,
            value_hidden_dim: dueling_hidden_dim,
            advantage_hidden_dim: dueling_hidden_dim,
            leaky_relu_alpha,
        }
    }
}

/// Dueling Q-Network with separate value and advantage streams
#[allow(missing_debug_implementations)]
pub struct DuelingQNetwork {
    /// Shared feature extraction layers
    shared_layers: Vec<Linear>,

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

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

    /// Configuration
    config: DuelingConfig,

    /// VarMap for weight management
    vars: VarMap,

    /// Device (CPU or CUDA)
    device: Device,
}

impl DuelingQNetwork {
    /// Create new dueling Q-network
    ///
    /// # Arguments
    ///
    /// * `config` - Dueling network configuration
    /// * `device` - Device to create network on (CPU or CUDA)
    ///
    /// # Returns
    ///
    /// New DuelingQNetwork instance with Xavier-initialized weights
    pub fn new(config: DuelingConfig, device: Device) -> Result<Self, MLError> {
        let vars = VarMap::new();
        let var_builder = VarBuilder::from_varmap(&vars, DType::F32, &device);

        // 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_name = format!("shared_{}", i);
            let layer_vb = var_builder.pp(&layer_name);
            let layer = linear_xavier(current_dim, hidden_dim, layer_vb).map_err(|e| {
                MLError::ModelError(format!("Failed to Xavier init shared layer {}: {}", i, e))
            })?;
            shared_layers.push(layer);
            current_dim = hidden_dim;
        }

        // Value stream
        let value_fc_vb = var_builder.pp("value_fc");
        let value_fc = linear_xavier(current_dim, config.value_hidden_dim, value_fc_vb)
            .map_err(|e| MLError::ModelError(format!("Failed to Xavier init value_fc: {}", e)))?;

        let value_out_vb = var_builder.pp("value_out");
        let value_out = linear_xavier(config.value_hidden_dim, 1, value_out_vb)
            .map_err(|e| MLError::ModelError(format!("Failed to Xavier init value_out: {}", e)))?;

        // Advantage stream
        let advantage_fc_vb = var_builder.pp("advantage_fc");
        let advantage_fc = linear_xavier(current_dim, config.advantage_hidden_dim, advantage_fc_vb)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to Xavier init advantage_fc: {}", e))
            })?;

        let advantage_out_vb = var_builder.pp("advantage_out");
        let advantage_out = linear_xavier(config.advantage_hidden_dim, config.num_actions, advantage_out_vb)
            .map_err(|e| {
                MLError::ModelError(format!("Failed to Xavier init advantage_out: {}", e))
            })?;

        Ok(Self {
            shared_layers,
            value_fc,
            value_out,
            advantage_fc,
            advantage_out,
            config,
            vars,
            device,
        })
    }

    /// Forward pass through dueling network
    ///
    /// # Arguments
    ///
    /// * `state` - State tensor [batch_size, state_dim]
    ///
    /// # Returns
    ///
    /// Q-values tensor [batch_size, num_actions]
    ///
    /// # Mathematical Formula
    ///
    /// Q(s,a) = V(s) + [A(s,a) - mean(A(s,·))]
    ///
    /// Where:
    /// - V(s): State value (scalar per sample)
    /// - A(s,a): Advantage per action
    /// - mean(A(s,·)): Mean advantage (ensures identifiability)
    pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
        // Shared feature extraction
        let mut h = state.clone();
        for (i, layer) in self.shared_layers.iter().enumerate() {
            h = layer.forward(&h).map_err(|e| {
                MLError::ModelError(format!("Shared layer {} forward failed: {}", i, e))
            })?;

            // LeakyReLU activation
            h = candle_nn::ops::leaky_relu(&h, self.config.leaky_relu_alpha).map_err(|e| {
                MLError::ModelError(format!("LeakyReLU failed at shared layer {}: {}", i, e))
            })?;
        }

        // Value stream: V(s) → [batch, 1]
        let v = self.value_fc.forward(&h).map_err(|e| {
            MLError::ModelError(format!("Value FC forward failed: {}", e))
        })?;
        let v = candle_nn::ops::leaky_relu(&v, self.config.leaky_relu_alpha).map_err(|e| {
            MLError::ModelError(format!("Value LeakyReLU failed: {}", e))
        })?;
        let v = self.value_out.forward(&v).map_err(|e| {
            MLError::ModelError(format!("Value output forward failed: {}", e))
        })?; // [batch, 1]

        // Advantage stream: A(s,a) → [batch, num_actions]
        let a = self.advantage_fc.forward(&h).map_err(|e| {
            MLError::ModelError(format!("Advantage FC forward failed: {}", e))
        })?;
        let a = candle_nn::ops::leaky_relu(&a, self.config.leaky_relu_alpha).map_err(|e| {
            MLError::ModelError(format!("Advantage LeakyReLU failed: {}", e))
        })?;
        let a = self.advantage_out.forward(&a).map_err(|e| {
            MLError::ModelError(format!("Advantage output forward failed: {}", e))
        })?; // [batch, num_actions]

        // Compute mean advantage: mean(A(s,·)) → [batch]
        // Use dimension 1 to average across actions (dim 0 is batch)
        let a_mean = a
            .mean(1)
            .map_err(|e| MLError::ModelError(format!("Advantage mean failed: {}", e)))?; // [batch]

        // Broadcast operations:
        // Q(s,a) = V(s) + A(s,a) - mean(A(s,·))
        //
        // Shapes:
        // - v: [batch, 1]
        // - a: [batch, num_actions]
        // - a_mean: [batch]
        //
        // Need to unsqueeze a_mean to [batch, 1] for broadcasting
        let a_mean_unsqueezed = a_mean.unsqueeze(1).map_err(|e| {
            MLError::ModelError(format!("Advantage mean unsqueeze failed: {}", e))
        })?; // [batch, 1]

        // Broadcast v to match advantage shape [batch, num_actions]
        let v_broadcast = v.broadcast_as(a.shape()).map_err(|e| {
            MLError::ModelError(format!("Value broadcast failed: {}", e))
        })?; // [batch, num_actions]

        // Broadcast a_mean to match advantage shape [batch, num_actions]
        let a_mean_broadcast = a_mean_unsqueezed.broadcast_as(a.shape()).map_err(|e| {
            MLError::ModelError(format!("Advantage mean broadcast failed: {}", e))
        })?; // [batch, num_actions]

        // Q = V + (A - mean(A))
        // All tensors now [batch, num_actions]
        let q_values = (&v_broadcast + &a - &a_mean_broadcast).map_err(|e| {
            MLError::ModelError(format!("Q-value combination failed: {}", e))
        })?;

        Ok(q_values)
    }

    /// Get VarMap for weight serialization
    pub fn vars(&self) -> &VarMap {
        &self.vars
    }

    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }

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

    /// Copy weights from another dueling network
    pub fn copy_weights_from(&mut self, other: &DuelingQNetwork) -> Result<(), MLError> {
        let self_vars = self.vars.data().lock().map_err(|e| {
            MLError::ConcurrencyError {
                operation: format!("lock self vars: {}", e),
            }
        })?;
        let other_vars = other.vars.data().lock().map_err(|e| {
            MLError::ConcurrencyError {
                operation: format!("lock other vars: {}", e),
            }
        })?;

        for (name, self_var) in self_vars.iter() {
            if let Some(other_var) = other_vars.get(name) {
                let other_tensor = other_var.as_tensor();
                self_var.set(other_tensor).map_err(|e| {
                    MLError::ModelError(format!("Failed to copy weight {}: {}", name, e))
                })?;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::Device;

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

        let device = Device::Cpu;
        let network = DuelingQNetwork::new(config, device)?;

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

    #[test]
    fn test_dueling_forward_pass() -> anyhow::Result<()> {
        let config = DuelingConfig::new(32, 45, vec![256, 128], 64, 64);
        let device = Device::Cpu;
        let network = DuelingQNetwork::new(config, device)?;

        // Create batch of states
        let batch_size = 4;
        let state = Tensor::randn(0f32, 1.0, (batch_size, 32), &Device::Cpu)?;

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

        // Check output shape
        assert_eq!(q_values.dims(), &[batch_size, 45]);

        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 device = Device::Cpu;
        let network = DuelingQNetwork::new(config, device)?;

        // Simple state
        let state = Tensor::ones((1, 4), DType::F32, &Device::Cpu)?;

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

        // Q-values should be valid (no NaN/Inf)
        let q_vec = q_values.to_vec2::<f32>()?;
        for &q in &q_vec[0] {
            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
            45,              // num_actions
            &[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, 45);
        assert_eq!(config.shared_hidden_dims, vec![256, 128]); // N-1 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 device = Device::Cpu;

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

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

        // Verify same output for same input
        let state = Tensor::ones((1, 8), DType::F32, &Device::Cpu)?;
        let q1 = network1.forward(&state)?;
        let q2 = network2.forward(&state)?;

        let q1_vec = q1.to_vec2::<f32>()?;
        let q2_vec = q2.to_vec2::<f32>()?;

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

        Ok(())
    }
}

// Manual Debug implementation for DuelingQNetwork (Wave 8.1 - Fix test compilation)
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())
            .field("device", &format!("{:?}", self.device))
            .finish()
    }
}