Skip to main content

rlox_rl_ops/
estimator.rs

1use crate::error::RlOpsError;
2
3/// Compute per-rollout advantages from a flat slice of scalar rewards.
4///
5/// `rewards` is a flat slice of length `n_groups * group_size`.
6/// Returns a `Vec` of the same length with per-rollout advantage estimates.
7///
8/// # Contract
9///
10/// Implementations MUST be:
11/// - **Deterministic**: same inputs always produce identical outputs.
12/// - **Free of autograd operations**: this crate has no tensor/gradient deps.
13/// - **Thread-safe**: `Send + Sync` so the estimator can live inside `Arc`.
14pub trait AdvantageEstimator: Send + Sync {
15    fn compute(&self, rewards: &[f32], group_size: usize) -> Result<Vec<f32>, RlOpsError>;
16}