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
24pub 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
36pub 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 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}