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

//! Continuous PPO Implementation
//!
//! This module provides a PPO implementation specifically designed for continuous
//! action spaces, using Gaussian policies for position sizing.

use candle_core::{DType, Device, Tensor};
use ml_core::mixed_precision::training_dtype;
use candle_nn::Optimizer; // Required for Adam::new and backward_step methods
use candle_optimisers::adam::Adam;
use candle_optimisers::adam::ParamsAdam;
use serde::{Deserialize, Serialize};
use tracing::debug;

use super::continuous_policy::ContinuousAction; // Keep only ContinuousAction
use super::flow_policy::{FlowPolicy, FlowPolicyConfig};
use super::gae::GAEConfig;
use super::ppo::ValueNetwork;
use ml_core::gradient_accumulation::clip_grads;
use ml_core::tensor_ops::TensorOps;
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>,
    log_probs: Vec<f32>,
    pub advantages: Vec<f32>,
    returns: Vec<f32>,
}

impl ContinuousTrajectoryBatch {
    /// Create batch from trajectories with computed advantages and returns
    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 for training stability
    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(())
    }

    /// Convert to tensors for training
    pub fn to_tensors(
        &self,
        device: &Device,
        state_dim: usize,
    ) -> Result<ContinuousTrajectoryTensors, MLError> {
        let batch_size = self.states.len();

        // Create state tensor
        let state_flat: Vec<f32> = self.states.iter().flatten().cloned().collect();
        let states = Tensor::from_vec(state_flat, (batch_size, state_dim), device)
            .map_err(|e| MLError::TrainingError(format!("Failed to create state tensor: {}", e)))?;

        // Create action tensor
        let actions =
            Tensor::from_vec(self.actions.clone(), (batch_size, 1), device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create action tensor: {}", e))
            })?;

        // Create other tensors
        let log_probs =
            Tensor::from_vec(self.log_probs.clone(), batch_size, device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create log_probs tensor: {}", e))
            })?;

        let advantages =
            Tensor::from_vec(self.advantages.clone(), batch_size, device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create advantages tensor: {}", e))
            })?;

        let returns = Tensor::from_vec(self.returns.clone(), batch_size, device).map_err(|e| {
            MLError::TrainingError(format!("Failed to create returns tensor: {}", e))
        })?;

        Ok(ContinuousTrajectoryTensors {
            states,
            actions,
            log_probs,
            advantages,
            returns,
        })
    }

    /// Create mini-batches for training
    pub fn create_mini_batches(&self, mini_batch_size: usize) -> Vec<ContinuousMiniBatch> {
        let mut mini_batches = Vec::new();
        let total_size = self.states.len();

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

            let mini_batch = ContinuousMiniBatch {
                states: self.states[start_idx..end_idx].to_vec(),
                actions: self.actions[start_idx..end_idx].to_vec(),
                log_probs: self.log_probs[start_idx..end_idx].to_vec(),
                advantages: self.advantages[start_idx..end_idx].to_vec(),
                returns: self.returns[start_idx..end_idx].to_vec(),
            };

            mini_batches.push(mini_batch);
        }

        mini_batches
    }
}

/// Mini-batch for continuous `PPO` training
#[derive(Debug, Clone)]
pub struct ContinuousMiniBatch {
    pub states: Vec<Vec<f32>>,
    actions: Vec<f32>,
    log_probs: Vec<f32>,
    advantages: Vec<f32>,
    returns: Vec<f32>,
}

impl ContinuousMiniBatch {
    /// Convert to tensors
    pub fn to_tensors(
        &self,
        device: &Device,
        state_dim: usize,
    ) -> Result<ContinuousTrajectoryTensors, MLError> {
        let batch_size = self.states.len();

        let state_flat: Vec<f32> = self.states.iter().flatten().cloned().collect();
        let states = Tensor::from_vec(state_flat, (batch_size, state_dim), device)
            .map_err(|e| MLError::TrainingError(format!("Failed to create state tensor: {}", e)))?;

        let actions =
            Tensor::from_vec(self.actions.clone(), (batch_size, 1), device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create action tensor: {}", e))
            })?;

        let log_probs =
            Tensor::from_vec(self.log_probs.clone(), batch_size, device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create log_probs tensor: {}", e))
            })?;

        let advantages =
            Tensor::from_vec(self.advantages.clone(), batch_size, device).map_err(|e| {
                MLError::TrainingError(format!("Failed to create advantages tensor: {}", e))
            })?;

        let returns = Tensor::from_vec(self.returns.clone(), batch_size, device).map_err(|e| {
            MLError::TrainingError(format!("Failed to create returns tensor: {}", e))
        })?;

        Ok(ContinuousTrajectoryTensors {
            states,
            actions,
            log_probs,
            advantages,
            returns,
        })
    }
}

/// Tensor representation of continuous trajectory batch
#[derive(Debug)]
pub struct ContinuousTrajectoryTensors {
    pub states: Tensor,
    pub actions: Tensor,
    pub log_probs: Tensor,
    pub advantages: Tensor,
    pub returns: Tensor,
}

/// Continuous `PPO` implementation for position sizing
#[allow(missing_debug_implementations)]
pub struct ContinuousPPO {
    /// Configuration
    config: ContinuousPPOConfig,
    /// Flow policy network (actor)
    pub actor: FlowPolicy,
    /// Value network (critic)
    pub critic: ValueNetwork,
    /// Policy optimizer
    policy_optimizer: Option<Adam>,
    /// Value optimizer
    value_optimizer: Option<Adam>,
    /// Training step counter
    training_steps: u64,
}

impl ContinuousPPO {
    /// Create new continuous `PPO`
    pub fn new(config: ContinuousPPOConfig) -> Result<Self, MLError> {
        let device = Device::Cpu; // Using CPU for compatibility

        // Ensure policy config has correct state dimension
        let mut flow_config = config.policy_config.clone();
        flow_config.state_dim = config.state_dim;
        flow_config.action_dim = 1;

        // Create actor network with FlowPolicy
        let actor = FlowPolicy::new(flow_config, &device)?;

        // Create critic network
        let critic = ValueNetwork::new(config.state_dim, &config.value_hidden_dims, device)?;

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

    /// Select action and get value estimate
    pub fn act(&self, state: &[f32]) -> Result<(ContinuousAction, f32), MLError> {
        let state_tensor = Tensor::from_vec(
            state.to_vec(),
            (1, self.config.state_dim),
            self.actor.device(),
        )
        .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?
        .to_dtype(DType::F32)?;

        // Get action from policy
        let (action_tensor, _log_prob) = self.actor.sample_action(&state_tensor)?;
        let action_value = action_tensor
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract action: {}", e)))?[0];
        let action = ContinuousAction::new(action_value);

        // Get value estimate
        let value = self
            .critic
            .forward(&state_tensor)?
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract value: {}", e)))?[0];

        Ok((action, value))
    }

    /// Get action with log probability (for trajectory collection)
    pub fn act_with_log_prob(
        &self,
        state: &[f32],
    ) -> Result<(ContinuousAction, f32, f32), MLError> {
        let state_tensor = Tensor::from_vec(
            state.to_vec(),
            (1, self.config.state_dim),
            self.actor.device(),
        )
        .map_err(|e| MLError::ModelError(format!("Failed to create state tensor: {}", e)))?
        .to_dtype(DType::F32)?;

        // Get action and log prob from policy
        let (action_tensor, log_prob_tensor) = self.actor.sample_action(&state_tensor)?;
        let action_value = action_tensor
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract action: {}", e)))?[0];
        let action = ContinuousAction::new(action_value);
        let log_prob = log_prob_tensor
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract log_prob: {}", e)))?[0];

        // Get value estimate
        let value = self
            .critic
            .forward(&state_tensor)?
            .flatten_all()?
            .to_vec1::<f32>()
            .map_err(|e| MLError::ModelError(format!("Failed to extract value: {}", e)))?[0];

        Ok((action, log_prob, value))
    }

    /// Update `PPO` networks with continuous trajectory batch
    pub fn update(&mut self, batch: &mut ContinuousTrajectoryBatch) -> Result<(f32, f32), MLError> {
        // Initialize optimizers if not done
        self.init_optimizers()?;

        // Normalize advantages
        batch.normalize_advantages()?;

        // Convert batch to tensors
        let device = self.actor.device();
        let _batch_tensors = batch.to_tensors(device, self.config.state_dim)?;

        let mut total_policy_loss = 0.0;
        let mut total_value_loss = 0.0;
        let mut num_updates = 0;

        // Train for multiple epochs
        for _epoch in 0..self.config.num_epochs {
            // Create mini-batches
            let mini_batches = batch.create_mini_batches(self.config.mini_batch_size);

            for mini_batch in mini_batches {
                let mini_tensors = mini_batch.to_tensors(device, self.config.state_dim)?;

                // Compute losses
                let policy_loss = self.compute_policy_loss(&mini_tensors)?;
                let value_loss = self.compute_value_loss(&mini_tensors)?;

                // Update policy network with gradient monitoring
                // Get vars before optimizer borrow to avoid borrow checker issues
                let actor_vars = self.actor.vars().all_vars();
                let grads = policy_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Policy backward failed: {}", e))
                })?;

                let mut grads = grads;
                let policy_grad_norm = clip_grads(
                    &mut grads,
                    &actor_vars,
                    self.config.max_grad_norm as f64,
                )?;

                debug!(
                    "Policy gradient norm: {:.4} (max: {:.4}, clipped: {})",
                    policy_grad_norm,
                    self.config.max_grad_norm,
                    policy_grad_norm > self.config.max_grad_norm as f64
                );

                if let Some(ref mut optimizer) = self.policy_optimizer {
                    optimizer.step(&grads).map_err(|e| {
                        MLError::TrainingError(format!("Policy optimizer step failed: {}", e))
                    })?;
                }

                // Update value network with gradient monitoring
                // Get vars before optimizer borrow to avoid borrow checker issues
                let critic_vars = self.critic.vars().all_vars();
                let value_grads = value_loss.backward().map_err(|e| {
                    MLError::TrainingError(format!("Value backward failed: {}", e))
                })?;

                let mut value_grads = value_grads;
                let value_grad_norm = clip_grads(
                    &mut value_grads,
                    &critic_vars,
                    self.config.max_grad_norm as f64,
                )?;

                debug!(
                    "Value gradient norm: {:.4} (max: {:.4}, clipped: {})",
                    value_grad_norm,
                    self.config.max_grad_norm,
                    value_grad_norm > self.config.max_grad_norm as f64
                );

                if let Some(ref mut optimizer) = self.value_optimizer {
                    optimizer.step(&value_grads).map_err(|e| {
                        MLError::TrainingError(format!("Value optimizer step failed: {}", e))
                    })?;
                }

                total_policy_loss += policy_loss.to_dtype(DType::F32).map_err(|e| {
                    MLError::TrainingError(format!("Failed to cast policy loss to F32: {}", e))
                })?.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract policy loss: {}", e))
                })?;
                total_value_loss += value_loss.to_dtype(DType::F32).map_err(|e| {
                    MLError::TrainingError(format!("Failed to cast value loss to F32: {}", e))
                })?.to_scalar::<f32>().map_err(|e| {
                    MLError::TrainingError(format!("Failed to extract value loss: {}", e))
                })?;
                num_updates += 1;
            }
        }

        self.training_steps += 1;

        let avg_policy_loss = total_policy_loss / num_updates as f32;
        let avg_value_loss = total_value_loss / num_updates as f32;

        Ok((avg_policy_loss, avg_value_loss))
    }

    /// Compute continuous `PPO` policy loss with clipping
    fn compute_policy_loss(&self, batch: &ContinuousTrajectoryTensors) -> Result<Tensor, MLError> {
        // Get current log determinants for flow-based actions (shape: [batch])
        let new_log_dets = self.actor.evaluate_actions(&batch.states, &batch.actions)?;

        // Compute probability ratio with clipping to prevent exp() overflow
        let log_ratio = (&new_log_dets - &batch.log_probs)?;

        // Clip log_ratio to [-20, 20] to prevent exp() overflow
        // exp(20) ≈ 4.85e8 (safe), exp(50) ≈ 5.18e21 (overflow to NaN)
        let log_ratio_min = Tensor::full(-20.0_f32, log_ratio.dims(), self.actor.device())
            .map_err(|e| MLError::TrainingError(format!("Failed to create log_ratio min tensor: {}", e)))?;
        let log_ratio_max = Tensor::full(20.0_f32, log_ratio.dims(), self.actor.device())
            .map_err(|e| MLError::TrainingError(format!("Failed to create log_ratio max tensor: {}", e)))?;
        let clipped_log_ratio = log_ratio.clamp(&log_ratio_min, &log_ratio_max)?;

        let ratio = clipped_log_ratio.exp()?;

        // Clipped surrogate objective
        let clip_epsilon_tensor = Tensor::from_vec(
            vec![self.config.clip_epsilon; batch.advantages.dims()[0]],
            batch.advantages.dims(),
            self.actor.device(),
        )
        .map_err(|e| MLError::TrainingError(format!("Failed to create clip tensor: {}", e)))?;

        let dtype = training_dtype(self.actor.device());
        let one_tensor = Tensor::ones(batch.advantages.dims(), dtype, self.actor.device())?;
        let clip_min = (&one_tensor - &clip_epsilon_tensor)?;
        let clip_max = (&one_tensor + &clip_epsilon_tensor)?;

        // Clamp ratio to [1-ε, 1+ε]
        let clipped_ratio = ratio.clamp(&clip_min, &clip_max)?;

        // PPO objective: min(ratio * advantage, clipped_ratio * advantage)
        let surr1 = (&ratio * &batch.advantages)?;
        let surr2 = (&clipped_ratio * &batch.advantages)?;
        let policy_loss_raw = TensorOps::elementwise_min(&surr1, &surr2)?;

        // Add entropy bonus for continuous actions
        let entropy = self.actor.entropy(&batch.states)?;
        let entropy_bonus = TensorOps::scalar_mul(&entropy, self.config.entropy_coeff as f64)?;

        // Final loss (negative because we want to maximize)
        let policy_loss_inner = (policy_loss_raw + entropy_bonus)?.mean_all()?;
        let policy_loss = TensorOps::negate(&policy_loss_inner)?;

        Ok(policy_loss)
    }

    /// Computes Huber loss for value function learning.
    ///
    /// Huber loss provides robust regression with continuous gradients:
    /// - Quadratic (MSE) for |error| <= delta: Smooth convergence near optimum
    /// - Linear for |error| > delta: Robust to outliers, bounded gradients
    ///
    /// **Mathematical Form**:
    /// ```
    /// L(x) = { 0.5 * x^2                if |x| <= delta
    ///        { delta * (|x| - 0.5*delta) if |x| > delta
    /// ```
    ///
    /// **Gradient Properties** (why this prevents vanishing gradients):
    /// - Quadratic region: ∂L/∂x = x  (bounded by ±delta)
    /// - Linear region: ∂L/∂x = ±delta (constant, non-zero)
    /// - **No dead zones**: Gradient always flows (unlike clamp where ∂clamp/∂x = 0)
    ///
    /// **Parameters**:
    /// - delta = 10.0: Transition threshold between quadratic and linear regions
    /// - Matches previous clamp threshold for consistency
    fn compute_value_loss(&self, batch: &ContinuousTrajectoryTensors) -> Result<Tensor, MLError> {
        let predicted_values = self.critic.forward(&batch.states)?;
        let value_diff = (&predicted_values - &batch.returns)?;

        // Huber loss: quadratic inside [-delta, delta], linear outside
        // Gradient is NEVER zero (prevents vanishing unlike clamp)
        let delta = 10.0_f32;
        let abs_diff = value_diff.abs()?;

        // Create delta tensor with same shape as abs_diff for broadcasting
        let dt = abs_diff.dtype();
        let delta_tensor = Tensor::full(delta, abs_diff.dims(), abs_diff.device())?.to_dtype(dt)?;
        let half_tensor = Tensor::full(0.5_f32, abs_diff.dims(), abs_diff.device())?.to_dtype(dt)?;
        let half_delta_sq = Tensor::full(0.5 * delta * delta, abs_diff.dims(), abs_diff.device())?.to_dtype(dt)?;

        // Mask: true if |value_diff| <= delta (quadratic region)
        let is_quadratic = abs_diff.le(&delta_tensor)?;

        // Quadratic loss: 0.5 * value_diff^2
        let quadratic_loss = value_diff.powf(2.0)?.mul(&half_tensor)?;

        // Linear loss: delta * (|value_diff| - 0.5 * delta)
        let linear_loss = abs_diff.mul(&delta_tensor)?.sub(&half_delta_sq)?;

        // Select based on mask
        let huber_loss = is_quadratic.where_cond(&quadratic_loss, &linear_loss)?
            .mean_all()?;

        let scaled_loss = TensorOps::scalar_mul(&huber_loss, self.config.value_loss_coeff as f64)?;

        Ok(scaled_loss)
    }

    /// Initialize optimizers
    fn init_optimizers(&mut self) -> Result<(), MLError> {
        if self.policy_optimizer.is_none() {
            let policy_params = ParamsAdam {
                lr: self.config.policy_learning_rate,
                beta_1: 0.9,
                beta_2: 0.999,
                eps: 1e-8,
                weight_decay: None,
                amsgrad: false,
            };
            self.policy_optimizer = Some(
                Adam::new(self.actor.vars().all_vars(), policy_params).map_err(|e| {
                    MLError::TrainingError(format!("Failed to create policy optimizer: {}", e))
                })?,
            );
        }

        if self.value_optimizer.is_none() {
            let value_params = ParamsAdam {
                lr: self.config.value_learning_rate,
                beta_1: 0.9,
                beta_2: 0.999,
                eps: 1e-8,
                weight_decay: None,
                amsgrad: false,
            };
            self.value_optimizer = Some(
                Adam::new(self.critic.vars().all_vars(), value_params).map_err(|e| {
                    MLError::TrainingError(format!("Failed to create value optimizer: {}", e))
                })?,
            );
        }

        Ok(())
    }

    /// 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 (log std)
    /// Note: Flows don't have a `log_std` parameter - returns 0.0 for compatibility
    pub const fn get_exploration_param(&self, _state: &[f32]) -> Result<f32, MLError> {
        Ok(0.0) // Flows don't have log_std
    }

    /// Set exploration parameter (for fixed std mode)
    /// Note: No-op for flows as they don't have a `log_std` parameter
    pub const fn set_exploration_param(&mut self, _log_std: f32) -> Result<(), MLError> {
        Ok(()) // No-op for flows
    }
}

/// 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 {
        // Get action and log probability
        let (action, log_prob, value) = agent.act_with_log_prob(&current_state)?;

        // Execute action in environment
        let (next_state, reward, done) = env_step_fn(&current_state, action)?;

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

        // Update state
        current_state = next_state;
        step_count += 1;

        if done {
            break;
        }
    }

    Ok(trajectory)
}

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

    #[test]
    fn test_continuous_ppo_creation() {
        let config = 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::default();
        let ppo = ContinuousPPO::new(config).unwrap();

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

        let (action, value) = result.unwrap();
        assert!(action.is_valid());
        assert!(action.position_size() >= 0.0 && action.position_size() <= 1.0);
        assert!(value.is_finite());
    }

    #[test]
    fn test_continuous_trajectory_step() {
        let action = ContinuousAction::new(0.5);
        let step = ContinuousTrajectoryStep::new(vec![0.1; 10], action, -1.5, 100.0, 50.0, false);

        assert_eq!(step.action.position_size(), 0.5);
        assert_eq!(step.reward, 100.0);
        assert!(!step.done);
    }

    #[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);

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

    #[test]
    fn test_tensor_conversion() {
        let action1 = ContinuousAction::new(0.4);
        let action2 = ContinuousAction::new(0.6);

        let step1 = ContinuousTrajectoryStep::new(vec![0.1; 8], action1, -1.2, 5.0, 2.5, false);

        let step2 = ContinuousTrajectoryStep::new(vec![0.2; 8], action2, -0.9, 15.0, 7.5, false);

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

        let trajectories = vec![trajectory];
        let advantages = vec![0.0, 0.0];
        let returns = vec![7.5, 22.5];

        let batch = ContinuousTrajectoryBatch::from_trajectories(trajectories, advantages, returns);

        let device = Device::Cpu;
        let tensors = batch.to_tensors(&device, 8);
        assert!(tensors.is_ok());

        let tensors = tensors.unwrap();
        assert_eq!(tensors.states.dims(), &[2, 8]);
        assert_eq!(tensors.actions.dims(), &[2, 1]);
        assert_eq!(tensors.advantages.dims(), &[2]);
    }

    #[test]
    fn test_exploration_parameter_control() {
        let config = ContinuousPPOConfig::default();
        let mut ppo = ContinuousPPO::new(config).unwrap();

        let state = vec![0.1; 64];

        // Get current exploration parameter
        let current_log_std = ppo.get_exploration_param(&state);
        assert!(current_log_std.is_ok());

        // Note: For flow policies, set_exploration_param is a no-op
        // Flow policies learn their exploration parameters through training
        let set_result = ppo.set_exploration_param(-2.0);
        // Should succeed as a no-op for flow policies
        assert!(set_result.is_ok());
    }
}