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. ```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; ``` Next, here is the core implementation of the flow matching objective. ```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 { // 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::()?; 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::()?; 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::()?; 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::()?; assert!((loss_val - (-1.2)).abs() < 1e-4); Ok(()) } } ```