foxhunt/crates/ml/zen_generated.code

<GENERATED-CODE>
I will create the core flow matching loss function in a new, self-contained module.

1.  Create a new directory `ml/src/ppo/flow_policy/` to house components related to the flow-based policy.
2.  Create a new module file `ml/src/ppo/flow_policy/mod.rs` to declare the `flow_matching` submodule.
3.  Create the implementation file `ml/src/ppo/flow_policy/flow_matching.rs`, which contains the PPO clipped surrogate objective adapted for flow-based policies.

<NEWFILE: ml/src/ppo/flow_policy/mod.rs>
```rust
//! # Flow-Based Policy Module
//!
//! This module contains components for implementing a flow-based policy within the PPO framework,
//! including the flow matching objective function.

pub mod flow_matching;
```
</NEWFILE>

Next, here is the core implementation of the flow matching objective.

<NEWFILE: ml/src/ppo/flow_policy/flow_matching.rs>
```rust
//! # Flow Matching Objective for Flow-Based Policy Optimization (FPO)
//!
//! This module implements the clipped surrogate objective for PPO using a flow-based policy.
//! Instead of using the ratio of likelihoods (log probabilities), it uses the ratio of
//! the determinants of the Jacobian of the flow transformation. This approach, detailed in
//! "Flow Matching Policy Optimization," helps mitigate gradient explosion issues common
//! with traditional policy gradient methods in continuous action spaces.
//!
//! The core idea is to replace the likelihood ratio `r(θ) = π_θ(a|s) / π_θ_old(a|s)` with
//! a flow ratio based on log-determinants: `r_flow = exp(log_det_new - log_det_old)`.
//! This ratio is then used in the standard PPO clipped surrogate objective.

use candle_core::{Device, Tensor};

use crate::{tensor_ops::TensorOps, MLError};

/// Configuration for the Flow Matching loss function.
#[derive(Debug, Clone, Copy)]
pub struct FlowMatchingConfig {
    /// The clipping parameter (epsilon) for the PPO surrogate objective.
    /// Typically set to 0.2.
    pub clip_epsilon: f32,
    /// The minimum value for clipping the log flow ratio to prevent underflow in `exp()`.
    pub log_ratio_clip_min: f32,
    /// The maximum value for clipping the log flow ratio to prevent overflow in `exp()`.
    pub log_ratio_clip_max: f32,
}

impl Default for FlowMatchingConfig {
    fn default() -> Self {
        Self {
            clip_epsilon: 0.2,
            log_ratio_clip_min: -20.0,
            log_ratio_clip_max: 20.0,
        }
    }
}

/// Computes the PPO clipped surrogate objective using the flow matching ratio.
///
/// This function is the core of the FPO update rule. It takes the log-determinants from the
/// current and old policies, computes the flow ratio, and then calculates the PPO loss.
///
/// # Arguments
/// * `new_log_dets` - A tensor of log-determinants from the current policy network for the actions in the batch. Shape: `[batch_size]`.
/// * `old_log_dets` - A tensor of log-determinants from the policy network used to collect the trajectory data. Shape: `[batch_size]`.
/// * `advantages` - A tensor of normalized advantages. Shape: `[batch_size]`.
/// * `config` - Configuration for the loss computation, including clipping parameters.
/// * `device` - The device on which to perform tensor operations.
///
/// # Returns
/// A scalar `Tensor` representing the final policy loss, ready for backpropagation.
/// The loss is negated because optimizers perform minimization, while PPO aims to maximize the objective.
pub fn compute_flow_matching_loss(
    new_log_dets: &Tensor,
    old_log_dets: &Tensor,
    advantages: &Tensor,
    config: &FlowMatchingConfig,
    device: &Device,
) -> Result<Tensor, MLError> {
    // 1. Compute the log of the flow ratio.
    // log_flow_ratio = log(det_new / det_old) = log_det_new - log_det_old
    let log_flow_ratio = (new_log_dets - old_log_dets)?;

    // 2. Clip the log ratio to prevent numerical instability (overflow/underflow) when exponentiating.
    // A range of [-20, 20] is a safe default, as exp(20) is large but manageable,
    // while exp(>~80) can lead to f32 infinity. This is a critical step for stability.
    let log_ratio_min = Tensor::full(config.log_ratio_clip_min, log_flow_ratio.dims(), device)
        .map_err(|e| MLError::TrainingError(format!("Failed to create log_ratio min tensor: {}", e)))?;
    let log_ratio_max = Tensor::full(config.log_ratio_clip_max, log_flow_ratio.dims(), device)
        .map_err(|e| MLError::TrainingError(format!("Failed to create log_ratio max tensor: {}", e)))?;
    let clipped_log_ratio = log_flow_ratio.clamp(&log_ratio_min, &log_ratio_max)?;

    // 3. Compute the flow ratio.
    let ratio = clipped_log_ratio.exp()?;

    // 4. Compute the clipped surrogate objective, as in standard PPO.
    // surrogate1 = ratio * advantage
    let surr1 = (&ratio * advantages)?;

    // surrogate2 = clamp(ratio, 1 - ε, 1 + ε) * advantage
    let clip_min = 1.0 - config.clip_epsilon;
    let clip_max = 1.0 + config.clip_epsilon;
    let clipped_ratio = ratio.clamp(clip_min, clip_max)?;
    let surr2 = (&clipped_ratio * advantages)?;

    // 5. The PPO objective is the minimum of the two surrogates.
    let policy_loss_raw = TensorOps::elementwise_min(&surr1, &surr2)?;

    // 6. The final loss is the negative mean of the objective function.
    // We negate it because we want to maximize the objective via gradient ascent,
    // which is equivalent to minimizing the negative objective.
    let policy_loss_mean = policy_loss_raw.mean_all()?;
    let final_loss = TensorOps::negate(&policy_loss_mean)?;

    Ok(final_loss)
}

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

    #[test]
    fn test_flow_matching_loss_computation() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = FlowMatchingConfig::default();

        // Batch size of 4
        let new_log_dets = Tensor::from_vec(vec![1.2, 0.8, -0.5, 2.0], 4, &device)?;
        let old_log_dets = Tensor::from_vec(vec![1.0, 1.0, -0.4, 1.5], 4, &device)?;
        let advantages = Tensor::from_vec(vec![1.5, -0.5, 2.0, 0.8], 4, &device)?;

        // Expected log_flow_ratio = [0.2, -0.2, -0.1, 0.5]
        // Expected ratio = [1.2214, 0.8187, 0.9048, 1.6487]

        // Expected clipped_ratio (epsilon=0.2, range=[0.8, 1.2])
        // clipped_ratio = [1.2, 0.8187, 0.9048, 1.2]

        // surr1 = [1.8321, -0.4093, 1.8096, 1.3189]
        // surr2 = [1.8, -0.4093, 1.8096, 0.96]

        // min(surr1, surr2) = [1.8, -0.4093, 1.8096, 0.96]
        // mean = (1.8 - 0.4093 + 1.8096 + 0.96) / 4 = 4.1603 / 4 = 1.040075
        // final_loss = -1.040075

        let loss =
            compute_flow_matching_loss(&new_log_dets, &old_log_dets, &advantages, &config, &device)?;
        let loss_val = loss.to_scalar::<f32>()?;

        assert!((loss_val - (-1.040075)).abs() < 1e-4);

        Ok(())
    }

    #[test]
    fn test_ratio_clipping() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = FlowMatchingConfig {
            clip_epsilon: 0.2,
            ..Default::default()
        };

        // This log_det difference will produce a large ratio that should be clipped
        let new_log_dets = Tensor::from_vec(vec![2.0], 1, &device)?;
        let old_log_dets = Tensor::from_vec(vec![0.0], 1, &device)?;
        let advantages = Tensor::from_vec(vec![10.0], 1, &device)?; // Positive advantage

        // log_flow_ratio = 2.0
        // ratio = exp(2.0) ≈ 7.389
        // clipped_ratio = clamp(7.389, 0.8, 1.2) = 1.2
        // surr1 = 7.389 * 10 ≈ 73.89
        // surr2 = 1.2 * 10 = 12.0
        // min = 12.0
        // loss = -12.0

        let loss =
            compute_flow_matching_loss(&new_log_dets, &old_log_dets, &advantages, &config, &device)?;
        let loss_val = loss.to_scalar::<f32>()?;
        assert!((loss_val - (-12.0)).abs() < 1e-4);

        // Test with negative advantage
        let advantages_neg = Tensor::from_vec(vec![-10.0], 1, &device)?;
        // surr1 = 7.389 * -10 ≈ -73.89
        // surr2 = 1.2 * -10 = -12.0
        // min(surr1, surr2) is surr1 because we want to decrease the probability of this action
        // min ≈ -73.89
        // loss = -(-73.89) ≈ 73.89
        let loss_neg_adv = compute_flow_matching_loss(
            &new_log_dets,
            &old_log_dets,
            &advantages_neg,
            &config,
            &device,
        )?;
        let loss_val_neg_adv = loss_neg_adv.to_scalar::<f32>()?;
        assert!((loss_val_neg_adv - 73.8905).abs() < 1e-4);

        Ok(())
    }

    #[test]
    fn test_log_ratio_clipping() -> Result<(), MLError> {
        let device = Device::Cpu;
        let config = FlowMatchingConfig {
            log_ratio_clip_min: -1.0,
            log_ratio_clip_max: 1.0,
            ..Default::default()
        };

        // This log_det difference is large and should be clipped before exp()
        let new_log_dets = Tensor::from_vec(vec![10.0], 1, &device)?;
        let old_log_dets = Tensor::from_vec(vec![0.0], 1, &device)?;
        let advantages = Tensor::from_vec(vec![1.0], 1, &device)?;

        // log_flow_ratio = 10.0
        // clipped_log_ratio = 1.0
        // ratio = exp(1.0) ≈ 2.718
        // clipped_ratio = clamp(2.718, 0.8, 1.2) = 1.2
        // surr1 ≈ 2.718
        // surr2 = 1.2
        // min = 1.2
        // loss = -1.2

        let loss =
            compute_flow_matching_loss(&new_log_dets, &old_log_dets, &advantages, &config, &device)?;
        let loss_val = loss.to_scalar::<f32>()?;
        assert!((loss_val - (-1.2)).abs() < 1e-4);

        Ok(())
    }
}
```
</GENERATED-CODE>