rlevo_evolution/
fitness.rs

1//! Fitness evaluation traits and adapters.
2//!
3//! Two traits model the two evaluation shapes strategies expect:
4//!
5//! - [`FitnessFn`] — evaluates a single member. Callers hand the fitness
6//!   function a host-side genome row (typically `Vec<f32>`) and receive a
7//!   scalar. Useful for simple benchmarks and for unit-testing operators.
8//! - [`BatchFitnessFn`] — evaluates an entire population in one call and
9//!   returns a device-resident `Tensor<B, 1>` of shape `(pop_size,)`. This
10//!   is the hot path — strategies call it once per generation.
11//!
12//! The [`FromFitnessEvaluable`] adapter bridges
13//! `rlevo-benchmarks::FitnessEvaluable<Individual = Vec<f64>, Landscape = L>`
14//! into a [`BatchFitnessFn`] over `Tensor<B, 2>`. It pulls each row to
15//! host, evaluates on the CPU, then stacks the results back onto the
16//! device. That is the straightforward (and only) path for fitness
17//! functions defined in terms of host-side scalar code. Purpose-built
18//! batched-on-device landscapes should implement [`BatchFitnessFn`]
19//! directly to avoid the round-trip.
20
21use burn::tensor::{Tensor, TensorData, backend::Backend};
22
23use rlevo_core::fitness::{FitnessEvaluable, Landscape};
24
25/// Single-member fitness evaluation.
26///
27/// Implementors may hold mutable state (e.g. a counter for number of
28/// evaluations) and are therefore `&mut self`.
29pub trait FitnessFn<G>: Send {
30    /// Evaluates one genome and returns its scalar fitness.
31    fn evaluate_one(&mut self, member: &G) -> f32;
32}
33
34/// Batched fitness evaluation over a population genome container `G`.
35///
36/// The returned tensor has shape `(pop_size,)` on the supplied device.
37/// Implementors must preserve row order — `fitness[i]` refers to the
38/// individual at row `i` of `population`.
39pub trait BatchFitnessFn<B: Backend, G>: Send {
40    /// Evaluates every member of `population` and stacks fitnesses.
41    fn evaluate_batch(&mut self, population: &G, device: &B::Device) -> Tensor<B, 1>;
42}
43
44/// Adapter from `FitnessEvaluable` to [`BatchFitnessFn<B, Tensor<B, 2>>`].
45///
46/// Each row of the population is pulled to host, converted to `Vec<f64>`,
47/// and passed to the underlying evaluator with the configured landscape.
48/// Fitness is computed on the host and then re-uploaded as a single
49/// `Tensor<B, 1>`.
50///
51/// # Precision
52///
53/// Populations are read as `f32` and widened to `f64` for the evaluator
54/// call; the returned `f64` fitness is narrowed back to `f32` before it
55/// is uploaded as a `Tensor<B, 1>`. Fitness values that exceed `f32`
56/// range (or rely on sub-ulp precision) will lose information at the
57/// narrowing step. Purpose-built batched-on-device landscapes should
58/// implement [`BatchFitnessFn`] directly to avoid the round-trip.
59///
60/// # Type Parameters
61///
62/// - `FE`: Concrete [`FitnessEvaluable`] implementation.
63/// - `L`: Landscape type; must match `FE::Landscape`.
64///
65/// # Panics
66///
67/// `evaluate_batch` panics if the supplied population tensor is not rank
68/// 2, or if its data cannot be read as `f32` (e.g. an integer backend).
69#[derive(Debug)]
70pub struct FromFitnessEvaluable<FE, L> {
71    evaluator: FE,
72    landscape: L,
73}
74
75impl<FE, L> FromFitnessEvaluable<FE, L> {
76    /// Builds the adapter from an evaluator and a landscape.
77    pub fn new(evaluator: FE, landscape: L) -> Self {
78        Self {
79            evaluator,
80            landscape,
81        }
82    }
83
84    /// Returns a reference to the wrapped landscape.
85    pub fn landscape(&self) -> &L {
86        &self.landscape
87    }
88}
89
90impl<FE, L, B> BatchFitnessFn<B, Tensor<B, 2>> for FromFitnessEvaluable<FE, L>
91where
92    B: Backend,
93    FE: FitnessEvaluable<Individual = Vec<f64>, Landscape = L> + Send,
94    L: Send + Sync,
95{
96    fn evaluate_batch(&mut self, population: &Tensor<B, 2>, device: &B::Device) -> Tensor<B, 1> {
97        let dims = population.shape().dims;
98        assert_eq!(dims.len(), 2, "population tensor must be rank 2");
99        let pop_size = dims[0];
100        let genome_dim = dims[1];
101
102        let flat = population
103            .clone()
104            .into_data()
105            .into_vec::<f32>()
106            .expect("tensor data must be readable as f32");
107        debug_assert_eq!(flat.len(), pop_size * genome_dim);
108
109        let mut fitness = Vec::with_capacity(pop_size);
110        let mut individual = Vec::with_capacity(genome_dim);
111        for row in 0..pop_size {
112            individual.clear();
113            let start = row * genome_dim;
114            individual.extend(
115                flat[start..start + genome_dim]
116                    .iter()
117                    .map(|&v| f64::from(v)),
118            );
119            let f = self.evaluator.evaluate(&individual, &self.landscape);
120            #[allow(clippy::cast_possible_truncation)]
121            fitness.push(f as f32);
122        }
123
124        let data = TensorData::new(fitness, [pop_size]);
125        Tensor::<B, 1>::from_data(data, device)
126    }
127}
128
129/// Adapter from [`Landscape`] to [`BatchFitnessFn<B, Tensor<B, 2>>`].
130///
131/// Use this when the landscape carries its own `evaluate(&[f64]) -> f64`
132/// (Sphere, Ackley, Rastrigin) so the example does not need a separate
133/// `FitnessEvaluable` shim. Each row is pulled to host as `f32`, widened
134/// to `f64`, evaluated, and re-uploaded as a `Tensor<B, 1>` — same
135/// precision caveats as [`FromFitnessEvaluable`] apply.
136///
137/// # Panics
138///
139/// `evaluate_batch` panics if the supplied population tensor is not rank
140/// 2, or if its data cannot be read as `f32` (e.g. an integer backend).
141#[derive(Debug)]
142pub struct FromLandscape<L> {
143    landscape: L,
144}
145
146impl<L> FromLandscape<L> {
147    /// Builds the adapter from a self-evaluating landscape.
148    pub fn new(landscape: L) -> Self {
149        Self { landscape }
150    }
151
152    /// Returns a reference to the wrapped landscape.
153    pub fn landscape(&self) -> &L {
154        &self.landscape
155    }
156}
157
158impl<L, B> BatchFitnessFn<B, Tensor<B, 2>> for FromLandscape<L>
159where
160    B: Backend,
161    L: Landscape,
162{
163    fn evaluate_batch(&mut self, population: &Tensor<B, 2>, device: &B::Device) -> Tensor<B, 1> {
164        let dims = population.shape().dims;
165        assert_eq!(dims.len(), 2, "population tensor must be rank 2");
166        let pop_size = dims[0];
167        let genome_dim = dims[1];
168
169        let flat = population
170            .clone()
171            .into_data()
172            .into_vec::<f32>()
173            .expect("tensor data must be readable as f32");
174        debug_assert_eq!(flat.len(), pop_size * genome_dim);
175
176        let mut fitness = Vec::with_capacity(pop_size);
177        let mut individual = Vec::with_capacity(genome_dim);
178        for row in 0..pop_size {
179            individual.clear();
180            let start = row * genome_dim;
181            individual.extend(
182                flat[start..start + genome_dim]
183                    .iter()
184                    .map(|&v| f64::from(v)),
185            );
186            let f = self.landscape.evaluate(&individual);
187            #[allow(clippy::cast_possible_truncation)]
188            fitness.push(f as f32);
189        }
190
191        let data = TensorData::new(fitness, [pop_size]);
192        Tensor::<B, 1>::from_data(data, device)
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use burn::backend::NdArray;
200    type TestBackend = NdArray;
201
202    #[derive(Debug, Clone, Copy)]
203    struct Sphere;
204
205    struct SphereFit;
206    impl FitnessEvaluable for SphereFit {
207        type Individual = Vec<f64>;
208        type Landscape = Sphere;
209        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
210            x.iter().map(|v| v * v).sum()
211        }
212    }
213
214    #[test]
215    fn from_fitness_evaluable_preserves_row_order() {
216        let device = Default::default();
217        let data = TensorData::new(
218            vec![1.0_f32, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0],
219            [3, 3],
220        );
221        let pop = Tensor::<TestBackend, 2>::from_data(data, &device);
222
223        let mut adapter = FromFitnessEvaluable::new(SphereFit, Sphere);
224        let fitness = adapter.evaluate_batch(&pop, &device);
225
226        let values = fitness.into_data().into_vec::<f32>().unwrap();
227        assert_eq!(values.len(), 3);
228        approx::assert_relative_eq!(values[0], 1.0, epsilon = 1e-6);
229        approx::assert_relative_eq!(values[1], 4.0, epsilon = 1e-6);
230        approx::assert_relative_eq!(values[2], 9.0, epsilon = 1e-6);
231    }
232
233    #[test]
234    fn from_landscape_preserves_row_order() {
235        struct SphereLandscape;
236        impl Landscape for SphereLandscape {
237            fn evaluate(&self, x: &[f64]) -> f64 {
238                x.iter().map(|v| v * v).sum()
239            }
240        }
241
242        let device = Default::default();
243        let data = TensorData::new(
244            vec![1.0_f32, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0],
245            [3, 3],
246        );
247        let pop = Tensor::<TestBackend, 2>::from_data(data, &device);
248
249        let mut adapter = FromLandscape::new(SphereLandscape);
250        let fitness = adapter.evaluate_batch(&pop, &device);
251
252        let values = fitness.into_data().into_vec::<f32>().unwrap();
253        assert_eq!(values.len(), 3);
254        approx::assert_relative_eq!(values[0], 1.0, epsilon = 1e-6);
255        approx::assert_relative_eq!(values[1], 4.0, epsilon = 1e-6);
256        approx::assert_relative_eq!(values[2], 9.0, epsilon = 1e-6);
257    }
258}
rlevo_evolution/fitness.rs

rlevo_evolution/
fitness.rs
