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///
18/// # Examples
19///
20/// ```
21/// use burn::backend::Flex;
22/// use burn::tensor::Tensor;
23/// use rlevo_evolution::shaping::z_score;
24///
25/// let device = Default::default();
26/// // Five fitness values: mean 3.0, all distinct.
27/// let t = Tensor::<Flex, 1>::from_floats([1.0f32, 2.0, 3.0, 4.0, 5.0], &device);
28/// let z = z_score(t);
29/// let values = z.into_data().into_vec::<f32>().unwrap();
30/// // After z-scoring the mean of the output is 0 (within floating-point tolerance).
31/// let mean: f32 = values.iter().sum::<f32>() / values.len() as f32;
32/// assert!(mean.abs() < 1e-5);
33/// ```
34#[must_use]
35pub fn z_score<B: Backend>(fitness: Tensor<B, 1>) -> Tensor<B, 1> {
36 let mean = fitness.clone().mean().into_scalar().elem::<f32>();
37 let n = fitness.dims()[0];
38 #[allow(clippy::cast_precision_loss)]
39 let n_f = n.max(1) as f32;
40 let centered = fitness - mean;
41 let var = centered.clone().powf_scalar(2.0).sum().into_scalar().elem::<f32>() / n_f;
42 let std = var.sqrt().max(1e-8);
43 centered / std
44}
45
46/// Returns centered ranks: the largest fitness maps to `+0.5`, the
47/// smallest to `-0.5`, with linear spacing in between.
48///
49/// Centered ranks are standard in modern ES (e.g. OpenAI-ES) because they
50/// remove outlier fitness magnitudes and keep the signal scale-free across
51/// generations. Implemented host-side because the argsort pathway is
52/// easier to reason about; swap in a tensor-native implementation if this
53/// ever shows up on a profile.
54///
55/// An empty input returns an empty tensor.
56///
57/// # Examples
58///
59/// ```
60/// use burn::backend::Flex;
61/// use burn::tensor::Tensor;
62/// use rlevo_evolution::shaping::centered_rank;
63///
64/// let device = Default::default();
65/// let t = Tensor::<Flex, 1>::from_floats([10.0f32, 20.0, 30.0, 40.0], &device);
66/// let r = centered_rank(t, &device);
67/// let values = r.into_data().into_vec::<f32>().unwrap();
68/// // Smallest value maps to -0.5, largest to +0.5.
69/// assert!((values[0] - (-0.5)).abs() < 1e-6);
70/// assert!((values[3] - 0.5).abs() < 1e-6);
71/// ```
72///
73/// # Panics
74///
75/// Panics if the tensor's element data cannot be converted to `f32` —
76/// for example, when using a backend that stores integer-typed tensors
77/// and `into_vec::<f32>()` returns an error.
78#[must_use]
79pub fn centered_rank<B: Backend>(fitness: Tensor<B, 1>, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> Tensor<B, 1> {
80 let data = fitness
81 .into_data()
82 .into_vec::<f32>()
83 .expect("centered_rank requires f32 tensor data");
84 let n = data.len();
85 if n == 0 {
86 return Tensor::<B, 1>::from_floats([0.0f32; 0], device);
87 }
88 let mut indices: Vec<usize> = (0..n).collect();
89 indices.sort_by(|&i, &j| data[i].partial_cmp(&data[j]).unwrap_or(std::cmp::Ordering::Equal));
90
91 #[allow(clippy::cast_precision_loss)]
92 let n_f = (n - 1).max(1) as f32;
93 let mut ranks = vec![0.0_f32; n];
94 for (rank, &idx) in indices.iter().enumerate() {
95 #[allow(clippy::cast_precision_loss)]
96 let r = rank as f32 / n_f - 0.5;
97 ranks[idx] = r;
98 }
99 Tensor::<B, 1>::from_floats(ranks.as_slice(), device)
100}
101
102#[cfg(test)]
103mod tests {
104 use super::*;
105 use burn::backend::Flex;
106 type TestBackend = Flex;
107
108 #[test]
109 #[allow(clippy::cast_precision_loss)]
110 fn z_score_zero_mean_unit_std() {
111 let device = Default::default();
112 let t = Tensor::<TestBackend, 1>::from_floats([1.0f32, 2.0, 3.0, 4.0, 5.0], &device);
113 let z = z_score(t);
114 let values = z.into_data().into_vec::<f32>().unwrap();
115 let mean: f32 = values.iter().sum::<f32>() / values.len() as f32;
116 approx::assert_relative_eq!(mean, 0.0, epsilon = 1e-5);
117 }
118
119 #[test]
120 fn centered_rank_spans_half_interval() {
121 let device = Default::default();
122 let t = Tensor::<TestBackend, 1>::from_floats([10.0f32, 20.0, 30.0, 40.0], &device);
123 let r = centered_rank(t, &device);
124 let values = r.into_data().into_vec::<f32>().unwrap();
125 // smallest → -0.5, largest → +0.5
126 approx::assert_relative_eq!(values[0], -0.5, epsilon = 1e-6);
127 approx::assert_relative_eq!(values[3], 0.5, epsilon = 1e-6);
128 }
129
130 #[test]
131 fn centered_rank_preserves_order() {
132 let device = Default::default();
133 let t = Tensor::<TestBackend, 1>::from_floats([3.0f32, 1.0, 2.0], &device);
134 let r = centered_rank(t, &device);
135 let values = r.into_data().into_vec::<f32>().unwrap();
136 // original: 3, 1, 2 → ranks sorted ascending: [1, 2, 3] at indices [1, 2, 0]
137 // rank-positions centered: index 1 → -0.5, index 2 → 0.0, index 0 → 0.5
138 approx::assert_relative_eq!(values[1], -0.5, epsilon = 1e-6);
139 approx::assert_relative_eq!(values[2], 0.0, epsilon = 1e-6);
140 approx::assert_relative_eq!(values[0], 0.5, epsilon = 1e-6);
141 }
142}