rlox_rl_ops/
grpo.rs

1//! Group-Relative Policy Optimisation (GRPO) advantage estimator.
2
3use crate::error::RlOpsError;
4use crate::estimator::AdvantageEstimator;
5use crate::kl::f32_ops::{compute_batch_group_advantages, compute_group_advantages};
6
7/// Group-Relative Policy Optimisation advantage estimator.
8///
9/// Normalises rewards within each group: `z = (r - mean) / std`.
10/// Returns zeros when `std < 1e-8` (constant-reward group).
11///
12/// This is the canonical GRPO implementation matching the math in the original
13/// `rlox-core/src/llm/ops.rs` — extracted here so `rlox-sandbox` can depend on
14/// this slim crate instead of all of `rlox-core`.
15pub struct GroupRelativeEstimator;
16
17impl AdvantageEstimator for GroupRelativeEstimator {
18    /// Compute per-rollout z-score normalised advantages.
19    ///
20    /// `rewards` must be a flat slice of length `n_groups * group_size`.
21    /// Groups are contiguous: `rewards[0..group_size]` is group 0, etc.
22    ///
23    /// Uses rayon parallel dispatch when `rewards.len() >= 4096` elements.
24    fn compute(&self, rewards: &[f32], group_size: usize) -> Result<Vec<f32>, RlOpsError> {
25        compute_batch_group_advantages(rewards, group_size)
26    }
27}
28
29/// Convenience free function: compute advantages for a single group.
30///
31/// Equivalent to calling `GroupRelativeEstimator.compute(rewards, rewards.len())`.
32/// Exported so callers that only have one group don't need to pass `group_size`.
33pub fn compute_single_group_advantages(rewards: &[f32]) -> Vec<f32> {
34    compute_group_advantages(rewards)
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    /// Core contract test for the `AdvantageEstimator` trait implementation.
42    ///
43    /// Input:  rewards = [1, 0, 1, 0], group_size = 4
44    /// Expected: advantages = [+1, -1, +1, -1] (z-score normalised)
45    ///
46    /// Verification:
47    ///   mean = 0.5, variance = 0.25, std = 0.5
48    ///   adv[i] = (r[i] - 0.5) / 0.5
49    ///          => (1 - 0.5)/0.5 = +1.0, (0 - 0.5)/0.5 = -1.0
50    #[test]
51    fn advantage_estimator_trait_binary_rewards() {
52        let estimator = GroupRelativeEstimator;
53        let rewards = [1.0f32, 0.0, 1.0, 0.0];
54        let adv = estimator
55            .compute(&rewards, 4)
56            .expect("compute must succeed");
57        assert_eq!(adv.len(), 4);
58
59        let expected = [1.0f32, -1.0, 1.0, -1.0];
60        for (i, (&got, &exp)) in adv.iter().zip(expected.iter()).enumerate() {
61            assert!(
62                (got - exp).abs() < 1e-5,
63                "adv[{i}]: expected {exp}, got {got}"
64            );
65        }
66    }
67
68    #[test]
69    fn advantage_estimator_trait_constant_group_returns_zeros() {
70        let estimator = GroupRelativeEstimator;
71        let rewards = [5.0f32, 5.0, 5.0, 5.0];
72        let adv = estimator
73            .compute(&rewards, 4)
74            .expect("compute must succeed");
75        assert!(
76            adv.iter().all(|&v| v == 0.0),
77            "constant rewards must produce all-zero advantages"
78        );
79    }
80
81    #[test]
82    fn advantage_estimator_trait_multiple_groups() {
83        let estimator = GroupRelativeEstimator;
84        // Two groups: [1,0,1,0] and [1,1,1,1]
85        let rewards = [1.0f32, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0];
86        let adv = estimator
87            .compute(&rewards, 4)
88            .expect("compute must succeed");
89        assert_eq!(adv.len(), 8);
90        // First group: z-scored [+1, -1, +1, -1]
91        assert!((adv[0] - 1.0).abs() < 1e-5);
92        assert!((adv[1] + 1.0).abs() < 1e-5);
93        // Second group: all same → all zeros
94        assert!(adv[4..8].iter().all(|&v| v == 0.0));
95    }
96
97    #[test]
98    fn advantage_estimator_trait_bad_group_size_returns_err() {
99        let estimator = GroupRelativeEstimator;
100        // 3 rewards with group_size=2 — not divisible
101        let result = estimator.compute(&[1.0f32, 0.0, 1.0], 2);
102        assert!(result.is_err(), "non-divisible group_size must return Err");
103
104        // group_size = 0 must return Err
105        let result2 = estimator.compute(&[1.0f32], 0);
106        assert!(result2.is_err(), "group_size=0 must return Err");
107    }
108
109    #[test]
110    fn single_group_convenience_function() {
111        let rewards = [2.0f32, 0.0];
112        let adv = compute_single_group_advantages(&rewards);
113        assert_eq!(adv.len(), 2);
114        assert!((adv[0] - 1.0).abs() < 1e-5);
115        assert!((adv[1] + 1.0).abs() < 1e-5);
116    }
117
118    /// Verify the estimator is object-safe and can be stored behind Arc<dyn>.
119    #[test]
120    fn advantage_estimator_is_dyn_compatible() {
121        use std::sync::Arc;
122        let estimator: Arc<dyn AdvantageEstimator> = Arc::new(GroupRelativeEstimator);
123        let rewards = [1.0f32, 0.0, 1.0, 0.0];
124        let adv = estimator.compute(&rewards, 4).unwrap();
125        assert_eq!(adv.len(), 4);
126    }
127}
rlox_rl_ops/grpo.rs

rlox_rl_ops/
grpo.rs
