foxhunt/crates/ml/src/ppo/trainable_adapter.rs

//! UnifiedTrainable implementation for PPO model
//!
//! This module provides the UnifiedTrainable trait implementation for PPO,
//! enabling standardized training orchestration with checkpoint management,
//! metrics collection, and batch training support.

use ml_core::device::MlDevice;
use ml_core::native_types::NativeDevice;
use std::collections::HashMap;
use std::path::Path;
use tracing::info;

use super::ppo::{PPOConfig, PPO};
use super::trajectories::TrajectoryBatch;
use crate::training::unified_trainer::{CheckpointMetadata, TrainingMetrics, UnifiedTrainable};
use crate::MLError;

/// Unified PPO model for standardized training
pub struct UnifiedPPO {
    /// Underlying PPO model
    ppo: PPO,
    /// Current training step
    step: usize,
    /// Policy learning rate
    policy_lr: f64,
    /// Value learning rate
    value_lr: f64,
    /// Last training loss (policy + value)
    last_policy_loss: f64,
    last_value_loss: f64,
    /// Gradient norm from last backward pass
    last_grad_norm: Option<f64>,
    /// Custom metrics storage
    custom_metrics: HashMap<String, f64>,
}

impl std::fmt::Debug for UnifiedPPO {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UnifiedPPO")
            .field("step", &self.step)
            .field("policy_lr", &self.policy_lr)
            .field("value_lr", &self.value_lr)
            .field("last_policy_loss", &self.last_policy_loss)
            .field("last_value_loss", &self.last_value_loss)
            .field("last_grad_norm", &self.last_grad_norm)
            .field("custom_metrics", &self.custom_metrics)
            .finish_non_exhaustive()
    }
}

impl UnifiedPPO {
    /// Create new UnifiedPPO from config and device
    pub fn new(config: PPOConfig, _device: MlDevice) -> Result<Self, MLError> {
        let policy_lr = config.policy_learning_rate;
        let value_lr = config.value_learning_rate;

        let ppo = PPO::new(config)?;

        Ok(Self {
            ppo,
            step: 0,
            policy_lr,
            value_lr,
            last_policy_loss: 0.0,
            last_value_loss: 0.0,
            last_grad_norm: None,
            custom_metrics: HashMap::new(),
        })
    }

    /// Get reference to underlying PPO model
    pub fn inner(&self) -> &PPO {
        &self.ppo
    }

    /// Get mutable reference to underlying PPO model
    pub fn inner_mut(&mut self) -> &mut PPO {
        &mut self.ppo
    }
}

impl UnifiedTrainable for UnifiedPPO {
    fn model_type(&self) -> &str {
        "PPO"
    }

    fn device_name(&self) -> String {
        NativeDevice::Cuda(0).to_string()
    }

    fn forward_loss(&mut self, _input: &[f32], _target: &[f32]) -> Result<f64, MLError> {
        // PPO uses trajectory-based update, not supervised forward-loss
        Ok(self.last_policy_loss + self.last_value_loss)
    }

    fn backward(&mut self, _loss_value: f64) -> Result<f64, MLError> {
        // PPO backward pass is integrated into update() method
        Ok(self.last_grad_norm.unwrap_or(0.0))
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        // PPO optimizer step is integrated into update() method
        Ok(())
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        // PPO uses Adam optimizer which handles gradient zeroing internally
        Ok(())
    }

    fn get_learning_rate(&self) -> f64 {
        self.policy_lr
    }

    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        self.policy_lr = lr;
        self.value_lr = lr;
        self.ppo.update_learning_rates(lr, lr)?;
        info!("Updated PPO learning rate to {}", lr);
        Ok(())
    }

    fn get_step(&self) -> usize {
        self.step
    }

    fn collect_metrics(&self) -> TrainingMetrics {
        let mut metrics = TrainingMetrics::default();

        metrics.loss = self.last_policy_loss + self.last_value_loss;
        metrics.learning_rate = self.policy_lr;
        metrics.grad_norm = self.last_grad_norm;

        metrics
            .custom_metrics
            .insert("policy_loss".to_owned(), self.last_policy_loss);
        metrics
            .custom_metrics
            .insert("value_loss".to_owned(), self.last_value_loss);
        metrics
            .custom_metrics
            .insert("policy_lr".to_owned(), self.policy_lr);
        metrics
            .custom_metrics
            .insert("value_lr".to_owned(), self.value_lr);

        for (key, value) in &self.custom_metrics {
            metrics.custom_metrics.insert(key.clone(), *value);
        }

        metrics
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        info!("Saving PPO checkpoint to {}", checkpoint_path);

        if let Some(parent) = Path::new(checkpoint_path).parent() {
            std::fs::create_dir_all(parent).map_err(|e| {
                MLError::CheckpointError(format!("Failed to create checkpoint directory: {}", e))
            })?;
        }

        // Save PPO checkpoint via the unified PPO::save_checkpoint API
        let ckpt_path = std::path::PathBuf::from(checkpoint_path);
        self.ppo.save_checkpoint(&ckpt_path).map_err(|e| {
            MLError::CheckpointError(format!("Failed to save PPO checkpoint: {}", e))
        })?;

        let metadata = CheckpointMetadata {
            model_type: "PPO".to_owned(),
            version: "1.0.0".to_owned(),
            epoch: self.step,
            step: self.step,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(self.ppo.get_config()).map_err(|e| {
                MLError::CheckpointError(format!("Failed to serialize config: {}", e))
            })?,
            metrics: self.collect_metrics(),
        };

        crate::training::unified_trainer::checkpoint::save_metadata(&metadata, checkpoint_path)?;

        info!("PPO checkpoint saved successfully");
        Ok(checkpoint_path.to_string())
    }

    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        info!("Loading PPO checkpoint from {}", checkpoint_path);

        let metadata =
            crate::training::unified_trainer::checkpoint::load_metadata(checkpoint_path)?;

        if metadata.model_type != "PPO" {
            return Err(MLError::CheckpointError(format!(
                "Model type mismatch: expected PPO, got {}",
                metadata.model_type
            )));
        }

        let config: PPOConfig = serde_json::from_value(metadata.config.clone()).map_err(|e| {
            MLError::CheckpointError(format!("Failed to deserialize config: {}", e))
        })?;

        let ckpt_path = std::path::PathBuf::from(checkpoint_path);
        self.ppo = PPO::load_checkpoint(&ckpt_path)?;

        self.step = metadata.step;
        self.policy_lr = config.policy_learning_rate;
        self.value_lr = config.value_learning_rate;

        info!("PPO checkpoint loaded successfully");
        Ok(metadata)
    }
}

/// Train PPO model using trajectory batch
pub fn train_batch_trajectory(
    unified_ppo: &mut UnifiedPPO,
    trajectory_batch: &mut TrajectoryBatch,
) -> Result<(f64, f64), MLError> {
    let (policy_loss, value_loss) = unified_ppo.inner_mut().update(trajectory_batch)?;

    unified_ppo.last_policy_loss = policy_loss as f64;
    unified_ppo.last_value_loss = value_loss as f64;
    unified_ppo.step += 1;
    unified_ppo.last_grad_norm = None;

    Ok((policy_loss as f64, value_loss as f64))
}

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

    #[test]
    fn test_unified_ppo_creation() -> Result<(), MLError> {
        let config = PPOConfig {
            state_dim: 16,
            num_actions: 63,
            policy_hidden_dims: vec![32],
            value_hidden_dims: vec![32],
            policy_learning_rate: 3e-4,
            value_learning_rate: 3e-4,
            ..Default::default()
        };

        let device = MlDevice::cuda(0)?;
        let ppo = UnifiedPPO::new(config, device)?;

        assert_eq!(ppo.model_type(), "PPO");
        assert_eq!(ppo.get_step(), 0);
        assert_eq!(ppo.get_learning_rate(), 3e-4);

        Ok(())
    }

    #[test]
    fn test_unified_ppo_metrics() -> Result<(), MLError> {
        let config = PPOConfig::default();
        let device = MlDevice::cuda(0)?;
        let ppo = UnifiedPPO::new(config, device)?;

        let metrics = ppo.collect_metrics();

        assert_eq!(metrics.learning_rate, ppo.get_learning_rate());
        assert!(metrics.custom_metrics.contains_key("policy_loss"));
        assert!(metrics.custom_metrics.contains_key("value_loss"));

        Ok(())
    }
}