radiate_core/
problem.rs

1use super::{Chromosome, Codec, Genotype, Score};
2use std::{
3    collections::VecDeque,
4    sync::{Arc, RwLock},
5};
6
7pub trait FitnessFunction<T, S = f32>: Send + Sync
8where
9    S: Into<Score>,
10{
11    fn evaluate(&self, individual: T) -> S;
12}
13
14impl<T, S, F> FitnessFunction<T, S> for F
15where
16    F: Fn(T) -> S + Send + Sync,
17    S: Into<Score>,
18{
19    fn evaluate(&self, individual: T) -> S {
20        (self)(individual)
21    }
22}
23
24/// [Problem] represents the interface for the fitness function or evaluation and encoding/decoding
25/// of a genotype to a phenotype within the genetic algorithm framework.
26///
27/// To run the genetic algorithm the three things that must be supplied are the encoding & decoding of
28/// the [Genotype] and the fitness function. [Problem] wraps all three into a
29/// single trait that can be supplied to the engine builder.
30pub trait Problem<C: Chromosome, T>: Send + Sync {
31    fn encode(&self) -> Genotype<C>;
32    fn decode(&self, genotype: &Genotype<C>) -> T;
33    fn eval(&self, individual: &Genotype<C>) -> Score;
34}
35
36/// [EngineProblem] is a generic, base level concrete implementation of the [Problem] trait that is the
37/// default problem used by the engine if none other is specified during building. We take the
38/// [Codec] and the fitness function from the user and simply wrap them into this struct.
39pub struct EngineProblem<C, T>
40where
41    C: Chromosome,
42{
43    pub codec: Arc<dyn Codec<C, T>>,
44    pub fitness_fn: Arc<dyn Fn(T) -> Score + Send + Sync>,
45}
46
47impl<C: Chromosome, T> Problem<C, T> for EngineProblem<C, T> {
48    fn encode(&self) -> Genotype<C> {
49        self.codec.encode()
50    }
51
52    fn decode(&self, genotype: &Genotype<C>) -> T {
53        self.codec.decode(genotype)
54    }
55
56    fn eval(&self, individual: &Genotype<C>) -> Score {
57        let phenotype = self.decode(individual);
58        (self.fitness_fn)(phenotype)
59    }
60}
61
62unsafe impl<C: Chromosome, T> Send for EngineProblem<C, T> {}
63unsafe impl<C: Chromosome, T> Sync for EngineProblem<C, T> {}
64
65pub struct CompositeFitnessFn<T, S> {
66    objectives: Vec<Arc<dyn for<'a> FitnessFunction<&'a T, S>>>,
67    weights: Vec<f32>,
68}
69
70impl<T, S> CompositeFitnessFn<T, S>
71where
72    S: Into<Score> + Clone,
73{
74    pub fn new() -> Self {
75        Self {
76            objectives: Vec::new(),
77            weights: Vec::new(),
78        }
79    }
80
81    pub fn add_weighted_fn(
82        mut self,
83        fitness_fn: impl for<'a> FitnessFunction<&'a T, S> + 'static,
84        weight: f32,
85    ) -> Self
86    where
87        S: Into<Score>,
88    {
89        self.objectives.push(Arc::new(fitness_fn));
90        self.weights.push(weight);
91        self
92    }
93
94    pub fn add_fitness_fn(
95        mut self,
96        fitness_fn: impl for<'a> FitnessFunction<&'a T, S> + 'static,
97    ) -> Self
98    where
99        S: Into<Score>,
100    {
101        self.objectives.push(Arc::new(fitness_fn));
102        self.weights.push(1.0);
103        self
104    }
105}
106
107impl<T> FitnessFunction<T> for CompositeFitnessFn<T, f32> {
108    fn evaluate(&self, individual: T) -> f32 {
109        let mut total_score = 0.0;
110        let mut total_weight = 0.0;
111        for (objective, weight) in self.objectives.iter().zip(&self.weights) {
112            let score = objective.evaluate(&individual);
113            total_score += score * weight;
114            total_weight += weight;
115        }
116
117        total_score / total_weight.max(1e-8)
118    }
119}
120
121#[derive(Clone)]
122pub struct NoveltySearch<T, BD>
123where
124    BD: Novelty<T> + Send + Sync,
125{
126    pub behavior: Arc<BD>,
127    pub archive: Arc<RwLock<VecDeque<BD::Descriptor>>>,
128    pub k: usize,
129    pub threshold: f32,
130    pub max_archive_size: usize,
131    __phantom: std::marker::PhantomData<T>,
132}
133
134impl<T, BD> NoveltySearch<T, BD>
135where
136    BD: Novelty<T> + Send + Sync,
137{
138    pub fn new(behavior: BD, k: usize, threshold: f32) -> Self {
139        NoveltySearch {
140            behavior: Arc::new(behavior),
141            archive: Arc::new(RwLock::new(VecDeque::new())),
142            k,
143            threshold,
144            max_archive_size: 1000,
145            __phantom: std::marker::PhantomData,
146        }
147    }
148
149    pub fn with_max_archive_size(mut self, size: usize) -> Self {
150        self.max_archive_size = size;
151        self
152    }
153
154    fn normalized_novelty_score(
155        &self,
156        descriptor: &BD::Descriptor,
157        archive: &VecDeque<BD::Descriptor>,
158    ) -> f32 {
159        if archive.is_empty() {
160            return 1.0;
161        }
162
163        let mut min_distance = f32::INFINITY;
164        let mut max_distance = f32::NEG_INFINITY;
165        let mut distances = archive
166            .iter()
167            .map(|archived| self.behavior.distance(descriptor, archived))
168            .inspect(|&d| {
169                max_distance = max_distance.max(d);
170                min_distance = min_distance.min(d);
171            })
172            .collect::<Vec<f32>>();
173
174        if max_distance == min_distance {
175            return 0.0;
176        }
177
178        distances.sort_by(|a, b| a.partial_cmp(b).unwrap());
179        let k = std::cmp::min(self.k, distances.len());
180        let k_nearest_distances = &distances[..k];
181        let avg_distance = k_nearest_distances.iter().sum::<f32>() / k as f32;
182
183        (avg_distance - min_distance) / (max_distance - min_distance)
184    }
185
186    fn evaluate_internal(&self, individual: &T) -> f32 {
187        let description = self.behavior.description(&individual);
188
189        let is_empty = {
190            let archive = self.archive.read().unwrap();
191            archive.is_empty()
192        };
193
194        if is_empty {
195            let mut writer = self.archive.write().unwrap();
196            writer.push_back(description);
197            return 1.0;
198        }
199
200        let (novelty, should_add) = {
201            let archive = self.archive.read().unwrap();
202            let result = self.normalized_novelty_score(&description, &archive);
203            let should_add = result > self.threshold || archive.len() < self.k;
204
205            (result, should_add)
206        };
207
208        let mut writer = self.archive.write().unwrap();
209
210        if should_add {
211            writer.push_back(description);
212            while writer.len() > self.max_archive_size {
213                writer.pop_front();
214            }
215        }
216
217        novelty
218    }
219}
220
221impl<T, BD> FitnessFunction<T, f32> for NoveltySearch<T, BD>
222where
223    BD: Novelty<T> + Send + Sync,
224    T: Send + Sync,
225{
226    fn evaluate(&self, individual: T) -> f32 {
227        self.evaluate_internal(&individual)
228    }
229}
230
231impl<T, BD> FitnessFunction<&T, f32> for NoveltySearch<T, BD>
232where
233    BD: Novelty<T> + Send + Sync,
234    T: Send + Sync,
235{
236    fn evaluate(&self, individual: &T) -> f32 {
237        self.evaluate_internal(individual)
238    }
239}
240
241pub trait Novelty<T>: Send + Sync {
242    type Descriptor: Send + Sync;
243
244    fn description(&self, phenotype: &T) -> Self::Descriptor;
245
246    fn distance(&self, a: &Self::Descriptor, b: &Self::Descriptor) -> f32;
247}
248
249pub struct FitnessDescriptor<F, T, S>
250where
251    F: for<'a> FitnessFunction<&'a T, S>,
252    S: Into<Score>,
253{
254    fitness_fn: Arc<F>,
255    _score_phantom: std::marker::PhantomData<S>,
256    _phantom: std::marker::PhantomData<T>,
257}
258
259impl<F, T, S> FitnessDescriptor<F, T, S>
260where
261    F: for<'a> FitnessFunction<&'a T, S> + 'static,
262    T: Send + Sync + 'static,
263    S: Into<Score> + Send + Sync,
264{
265    /// Create a new fitness descriptor that uses the output of a fitness function as the behavioral descriptor.
266    /// This allows you to use fitness scores directly as behavioral descriptors for novelty search or diversity measurement.
267    pub fn new(fitness_fn: F) -> Self {
268        Self {
269            fitness_fn: Arc::new(fitness_fn),
270            _score_phantom: std::marker::PhantomData,
271            _phantom: std::marker::PhantomData,
272        }
273    }
274}
275
276impl<F, T, S> Novelty<T> for FitnessDescriptor<F, T, S>
277where
278    F: for<'a> FitnessFunction<&'a T, S> + 'static,
279    T: Send + Sync + 'static,
280    S: Into<Score> + Send + Sync,
281{
282    type Descriptor = Vec<f32>;
283
284    fn description(&self, phenotype: &T) -> Self::Descriptor {
285        let score = self.fitness_fn.evaluate(phenotype);
286        score.into().into()
287    }
288
289    fn distance(&self, a: &Self::Descriptor, b: &Self::Descriptor) -> f32 {
290        if a.len() != b.len() {
291            return f32::INFINITY;
292        }
293
294        a.iter()
295            .zip(b.iter())
296            .map(|(x, y)| (x - y).powi(2))
297            .sum::<f32>()
298            .sqrt()
299    }
300}