foxhunt/crates/ml-ppo/src/continuous_ppo.rs

//! Continuous PPO Implementation
//!
//! PPO for continuous action spaces using flow-based policies for position sizing.
//! All computation is GPU-native via `cuda_nn` primitives.

use serde::{Deserialize, Serialize};
use super::continuous_policy::ContinuousAction;
use super::flow_policy::{FlowPolicy, FlowPolicyConfig};
use super::gae::GAEConfig;
use ml_core::MLError;

/// Configuration for Continuous `PPO`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContinuousPPOConfig {
    /// State dimension
    pub state_dim: usize,
    /// Flow policy configuration
    pub policy_config: FlowPolicyConfig,
    /// Value network hidden dimensions
    pub value_hidden_dims: Vec<usize>,
    /// Learning rates
    pub policy_learning_rate: f64,
    pub value_learning_rate: f64,
    /// `PPO` clip parameter (epsilon)
    pub clip_epsilon: f32,
    /// Value function loss coefficient
    pub value_loss_coeff: f32,
    /// Entropy coefficient for exploration
    pub entropy_coeff: f32,
    /// GAE configuration
    pub gae_config: GAEConfig,
    /// Training parameters
    pub batch_size: usize,
    pub mini_batch_size: usize,
    pub num_epochs: usize,
    /// Maximum gradient norm for clipping
    pub max_grad_norm: f32,
}

impl Default for ContinuousPPOConfig {
    fn default() -> Self {
        Self {
            state_dim: 64,
            policy_config: FlowPolicyConfig::default(),
            value_hidden_dims: vec![128, 64],
            policy_learning_rate: 3e-4,
            value_learning_rate: 3e-4,
            clip_epsilon: 0.2,
            value_loss_coeff: 0.5,
            entropy_coeff: 0.01,
            gae_config: GAEConfig::default(),
            batch_size: 2048,
            mini_batch_size: 64,
            num_epochs: 10,
            max_grad_norm: 0.5,
        }
    }
}

/// Continuous trajectory step for position sizing
#[derive(Debug, Clone)]
pub struct ContinuousTrajectoryStep {
    /// State observation
    pub state: Vec<f32>,
    /// Continuous action taken
    pub action: ContinuousAction,
    /// Log probability of the action
    pub log_prob: f32,
    /// Reward received
    pub reward: f32,
    /// Value estimate at this state
    pub value: f32,
    /// Whether this step terminated the episode
    pub done: bool,
}

impl ContinuousTrajectoryStep {
    pub const fn new(
        state: Vec<f32>,
        action: ContinuousAction,
        log_prob: f32,
        reward: f32,
        value: f32,
        done: bool,
    ) -> Self {
        Self { state, action, log_prob, reward, value, done }
    }
}

/// Continuous trajectory for collecting experiences
#[derive(Debug, Clone)]
pub struct ContinuousTrajectory {
    steps: Vec<ContinuousTrajectoryStep>,
}

impl Default for ContinuousTrajectory {
    fn default() -> Self {
        Self::new()
    }
}

impl ContinuousTrajectory {
    pub const fn new() -> Self {
        Self { steps: Vec::new() }
    }

    pub fn add_step(&mut self, step: ContinuousTrajectoryStep) {
        self.steps.push(step);
    }

    pub fn steps(&self) -> &[ContinuousTrajectoryStep] {
        &self.steps
    }

    pub fn len(&self) -> usize {
        self.steps.len()
    }

    pub fn is_empty(&self) -> bool {
        self.steps.is_empty()
    }
}

/// Batch of continuous trajectories for training
#[derive(Debug, Clone)]
pub struct ContinuousTrajectoryBatch {
    states: Vec<Vec<f32>>,
    pub actions: Vec<f32>,
    #[allow(dead_code)]
    log_probs: Vec<f32>,
    pub advantages: Vec<f32>,
    #[allow(dead_code)]
    returns: Vec<f32>,
}

impl ContinuousTrajectoryBatch {
    /// Create batch from trajectories
    pub fn from_trajectories(
        trajectories: Vec<ContinuousTrajectory>,
        advantages: Vec<f32>,
        returns: Vec<f32>,
    ) -> Self {
        let mut states = Vec::new();
        let mut actions = Vec::new();
        let mut log_probs = Vec::new();

        for trajectory in trajectories {
            for step in trajectory.steps() {
                states.push(step.state.clone());
                actions.push(step.action.position_size());
                log_probs.push(step.log_prob);
            }
        }

        Self { states, actions, log_probs, advantages, returns }
    }

    /// Normalize advantages
    pub fn normalize_advantages(&mut self) -> Result<(), MLError> {
        if self.advantages.is_empty() {
            return Ok(());
        }

        let mean: f32 = self.advantages.iter().sum::<f32>() / self.advantages.len() as f32;
        let variance: f32 = self.advantages.iter()
            .map(|&x| (x - mean).powi(2))
            .sum::<f32>() / self.advantages.len() as f32;
        let std = (variance + 1e-8).sqrt();

        for advantage in &mut self.advantages {
            *advantage = (*advantage - mean) / std;
        }

        Ok(())
    }

    /// Create mini-batch ranges for GPU-side `DtoD` slicing.
    ///
    /// Returns `(start, end)` index pairs. Callers should upload the full
    /// batch to GPU once and use range-based sub-batch extraction.
    pub fn create_mini_batch_ranges(&self) -> Vec<(usize, usize)> {
        // This method exists for API symmetry; callers can also iterate
        // manually. The key invariant: NO `.to_vec()` copies.
        let total_size = self.states.len();
        let mini_batch_size = total_size.max(1); // default full batch
        let mut ranges = Vec::new();

        for start_idx in (0..total_size).step_by(mini_batch_size) {
            let end_idx = (start_idx + mini_batch_size).min(total_size);
            ranges.push((start_idx, end_idx));
        }

        ranges
    }
}

/// Continuous PPO implementation for position sizing.
///
/// Uses flow-based policy for the actor and `CudaValueNetwork` for the critic.
#[allow(missing_debug_implementations)]
pub struct ContinuousPPO {
    /// Configuration
    config: ContinuousPPOConfig,
    /// Flow policy network (actor)
    pub actor: FlowPolicy,
    /// Value network (critic) - using `CudaValueNetwork`
    pub critic: crate::cuda_nn::CudaValueNetwork,
    /// Training step counter
    training_steps: u64,
}

impl ContinuousPPO {
    /// Create new continuous PPO
    pub fn new(config: ContinuousPPOConfig) -> Result<Self, MLError> {
        let mut flow_config = config.policy_config.clone();
        flow_config.state_dim = config.state_dim;
        flow_config.action_dim = 1;

        let actor = FlowPolicy::new(flow_config)?;

        let gpu_ctx = crate::cuda_nn::GpuContext::new()?;
        let critic = crate::cuda_nn::CudaValueNetwork::new(
            config.state_dim,
            &config.value_hidden_dims,
            gpu_ctx,
        )?;

        Ok(Self {
            config,
            actor,
            critic,
            training_steps: 0,
        })
    }

    /// Select action and get value estimate
    pub fn act(&self, state: &[f32]) -> Result<(ContinuousAction, f32), MLError> {
        let (action_vec, _log_prob) = self.actor.sample_action_host(state)?;
        let action_value = action_vec.first().copied().unwrap_or(0.5);
        let action = ContinuousAction::new(action_value);

        // Value from critic
        let gpu_ctx = self.critic.ctx();
        let input = crate::cuda_nn::networks::states_to_gpu(gpu_ctx, state)?;
        let value_gpu = self.critic.forward(&input.data, 1)?;
        let value_host = crate::cuda_nn::networks::gpu_to_host(gpu_ctx, &value_gpu)?;
        let value = value_host.first().copied().unwrap_or(0.0);

        Ok((action, value))
    }

    /// Get action with log probability
    pub fn act_with_log_prob(&self, state: &[f32]) -> Result<(ContinuousAction, f32, f32), MLError> {
        let (action_vec, log_prob_vec) = self.actor.sample_action_host(state)?;
        let action_value = action_vec.first().copied().unwrap_or(0.5);
        let action = ContinuousAction::new(action_value);
        let log_prob = log_prob_vec.first().copied().unwrap_or(0.0);

        let gpu_ctx = self.critic.ctx();
        let input = crate::cuda_nn::networks::states_to_gpu(gpu_ctx, state)?;
        let value_gpu = self.critic.forward(&input.data, 1)?;
        let value_host = crate::cuda_nn::networks::gpu_to_host(gpu_ctx, &value_gpu)?;
        let value = value_host.first().copied().unwrap_or(0.0);

        Ok((action, log_prob, value))
    }

    /// Get training steps
    pub const fn get_training_steps(&self) -> u64 {
        self.training_steps
    }

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

    /// Get current exploration parameter (no-op for flows)
    pub const fn get_exploration_param(&self, _state: &[f32]) -> Result<f32, MLError> {
        Ok(0.0)
    }

    /// Set exploration parameter (no-op for flows)
    pub const fn set_exploration_param(&mut self, _log_std: f32) -> Result<(), MLError> {
        Ok(())
    }
}

/// Utility function to collect continuous trajectories
pub fn collect_continuous_trajectories<F>(
    agent: &ContinuousPPO,
    env_step_fn: F,
    initial_state: Vec<f32>,
    max_steps: usize,
) -> Result<ContinuousTrajectory, MLError>
where
    F: Fn(&[f32], ContinuousAction) -> Result<(Vec<f32>, f32, bool), MLError>,
{
    let mut trajectory = ContinuousTrajectory::new();
    let mut current_state = initial_state;
    let mut step_count = 0;

    while step_count < max_steps {
        let (action, log_prob, value) = agent.act_with_log_prob(&current_state)?;
        let (next_state, reward, done) = env_step_fn(&current_state, action)?;

        trajectory.add_step(ContinuousTrajectoryStep::new(
            current_state.clone(),
            action,
            log_prob,
            reward,
            value,
            done,
        ));

        current_state = next_state;
        step_count += 1;

        if done {
            break;
        }
    }

    Ok(trajectory)
}

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

    #[test]
    fn test_continuous_ppo_creation() {
        let config = ContinuousPPOConfig {
            state_dim: 8,
            policy_config: FlowPolicyConfig {
                state_dim: 8,
                action_dim: 1,
                context_dim: 16,
                num_layers: 2,
                scale_clamp: 5.0,
            },
            value_hidden_dims: vec![16, 8],
            ..ContinuousPPOConfig::default()
        };
        let ppo = ContinuousPPO::new(config);
        assert!(ppo.is_ok());

        let ppo = ppo.unwrap();
        assert_eq!(ppo.get_training_steps(), 0);
    }

    #[test]
    fn test_continuous_action_selection() {
        let config = ContinuousPPOConfig {
            state_dim: 8,
            policy_config: FlowPolicyConfig {
                state_dim: 8,
                action_dim: 1,
                context_dim: 16,
                num_layers: 2,
                scale_clamp: 5.0,
            },
            value_hidden_dims: vec![16, 8],
            ..ContinuousPPOConfig::default()
        };
        let ppo = ContinuousPPO::new(config).unwrap();

        let state = vec![0.1; 8];
        let result = ppo.act(&state);
        assert!(result.is_ok());

        let (action, value) = result.unwrap();
        assert!(action.is_valid());
        assert!(value.is_finite());
    }

    #[test]
    fn test_continuous_trajectory_batch() {
        let action1 = ContinuousAction::new(0.3);
        let action2 = ContinuousAction::new(0.7);

        let step1 = ContinuousTrajectoryStep::new(vec![0.1; 4], action1, -1.0, 10.0, 5.0, false);
        let step2 = ContinuousTrajectoryStep::new(vec![0.2; 4], action2, -0.8, 20.0, 15.0, true);

        let mut trajectory = ContinuousTrajectory::new();
        trajectory.add_step(step1);
        trajectory.add_step(step2);

        let trajectories = vec![trajectory];
        let advantages = vec![0.1, 0.2];
        let returns = vec![15.0, 35.0];

        let mut batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);
        assert_eq!(batch.actions.len(), 2);
        assert_eq!(batch.states.len(), 2);

        let result = batch.normalize_advantages();
        assert!(result.is_ok());
    }
}