diff --git a/crates/ml/src/ppo/continuous_policy.rs b/crates/ml/src/ppo/continuous_policy.rs index 84e8f586b..909b64eea 100644 --- a/crates/ml/src/ppo/continuous_policy.rs +++ b/crates/ml/src/ppo/continuous_policy.rs @@ -330,7 +330,7 @@ impl ContinuousPolicyNetwork { let log_prob = (log_2pi_tensor * (-0.5))? .sub(&log_stds)? - .sub(&(squared_diff * (-0.5))?)?; + .sub(&(squared_diff * 0.5)?)?; Ok(log_prob.squeeze(1)?) // Remove extra dimension if present } @@ -567,9 +567,8 @@ mod tests { // Action should be in bounds assert!(action >= 0.0 && action <= 1.0); - // Log prob should be finite and negative + // Log prob should be finite (can be positive for continuous distributions with small σ) assert!(log_prob.is_finite()); - assert!(log_prob <= 0.0); } #[test] @@ -592,7 +591,7 @@ mod tests { assert_eq!(log_probs.dims(), &[2]); let log_probs_vec = log_probs.to_vec1::().unwrap(); - assert!(log_probs_vec.iter().all(|&lp| lp.is_finite() && lp <= 0.0)); + assert!(log_probs_vec.iter().all(|&lp| lp.is_finite())); } #[test]