rlevo_evolution/
shaping.rs

1//! Fitness shaping transforms.
2//!
3//! Strategies occasionally benefit from monotone transforms of raw
4//! fitness values: centered-rank mapping flattens outliers, z-scoring
5//! normalizes scale across generations, and weight decay penalizes
6//! large-norm genomes (common in neuroevolution). These helpers work
7//! directly on device tensors and are pure functions (RNG-free).
8
9use burn::prelude::ElementConversion;
10use burn::tensor::{backend::Backend, Tensor};
11
12/// Returns `fitness - fitness.mean()` divided by the (population) std-dev,
13/// clamped to avoid divide-by-zero when all fitnesses are equal.
14///
15/// The std-dev floor is `1e-8`; degenerate populations (all-equal fitness)
16/// therefore map to a vector of zeros rather than producing NaNs.
17#[must_use]
18pub fn z_score<B: Backend>(fitness: Tensor<B, 1>) -> Tensor<B, 1> {
19    let mean = fitness.clone().mean().into_scalar().elem::<f32>();
20    let n = fitness.shape().dims[0];
21    #[allow(clippy::cast_precision_loss)]
22    let n_f = n.max(1) as f32;
23    let centered = fitness - mean;
24    let var = centered.clone().powf_scalar(2.0).sum().into_scalar().elem::<f32>() / n_f;
25    let std = var.sqrt().max(1e-8);
26    centered / std
27}
28
29/// Returns centered ranks: the largest fitness maps to `+0.5`, the
30/// smallest to `-0.5`, with linear spacing in between.
31///
32/// Centered ranks are standard in modern ES (e.g. OpenAI-ES) because they
33/// remove outlier fitness magnitudes and keep the signal scale-free across
34/// generations. Implemented host-side because the argsort pathway is
35/// easier to reason about; swap in a tensor-native implementation if this
36/// ever shows up on a profile.
37///
38/// An empty input returns an empty tensor.
39///
40/// # Panics
41///
42/// Panics if `fitness`'s data cannot be read as `f32` (e.g. an integer
43/// backend tensor was passed in).
44#[must_use]
45pub fn centered_rank<B: Backend>(fitness: Tensor<B, 1>, device: &B::Device) -> Tensor<B, 1> {
46    let data = fitness
47        .into_data()
48        .into_vec::<f32>()
49        .expect("centered_rank requires f32 tensor data");
50    let n = data.len();
51    if n == 0 {
52        return Tensor::<B, 1>::from_floats([0.0f32; 0], device);
53    }
54    let mut indices: Vec<usize> = (0..n).collect();
55    indices.sort_by(|&i, &j| data[i].partial_cmp(&data[j]).unwrap_or(std::cmp::Ordering::Equal));
56
57    #[allow(clippy::cast_precision_loss)]
58    let n_f = (n - 1).max(1) as f32;
59    let mut ranks = vec![0.0_f32; n];
60    for (rank, &idx) in indices.iter().enumerate() {
61        #[allow(clippy::cast_precision_loss)]
62        let r = rank as f32 / n_f - 0.5;
63        ranks[idx] = r;
64    }
65    Tensor::<B, 1>::from_floats(ranks.as_slice(), device)
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use burn::backend::NdArray;
72    type TestBackend = NdArray;
73
74    #[test]
75    #[allow(clippy::cast_precision_loss)]
76    fn z_score_zero_mean_unit_std() {
77        let device = Default::default();
78        let t = Tensor::<TestBackend, 1>::from_floats([1.0f32, 2.0, 3.0, 4.0, 5.0], &device);
79        let z = z_score(t);
80        let values = z.into_data().into_vec::<f32>().unwrap();
81        let mean: f32 = values.iter().sum::<f32>() / values.len() as f32;
82        approx::assert_relative_eq!(mean, 0.0, epsilon = 1e-5);
83    }
84
85    #[test]
86    fn centered_rank_spans_half_interval() {
87        let device = Default::default();
88        let t = Tensor::<TestBackend, 1>::from_floats([10.0f32, 20.0, 30.0, 40.0], &device);
89        let r = centered_rank(t, &device);
90        let values = r.into_data().into_vec::<f32>().unwrap();
91        // smallest → -0.5, largest → +0.5
92        approx::assert_relative_eq!(values[0], -0.5, epsilon = 1e-6);
93        approx::assert_relative_eq!(values[3], 0.5, epsilon = 1e-6);
94    }
95
96    #[test]
97    fn centered_rank_preserves_order() {
98        let device = Default::default();
99        let t = Tensor::<TestBackend, 1>::from_floats([3.0f32, 1.0, 2.0], &device);
100        let r = centered_rank(t, &device);
101        let values = r.into_data().into_vec::<f32>().unwrap();
102        // original: 3, 1, 2 → ranks sorted ascending: [1, 2, 3] at indices [1, 2, 0]
103        // rank-positions centered: index 1 → -0.5, index 2 → 0.0, index 0 → 0.5
104        approx::assert_relative_eq!(values[1], -0.5, epsilon = 1e-6);
105        approx::assert_relative_eq!(values[2], 0.0, epsilon = 1e-6);
106        approx::assert_relative_eq!(values[0], 0.5, epsilon = 1e-6);
107    }
108}
rlevo_evolution/shaping.rs

rlevo_evolution/
shaping.rs
