1use std::marker::PhantomData;
24
25use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
26use rand::Rng;
27use rand::RngExt;
28
29use rlevo_core::bounds::Bounds;
30use rlevo_core::config::{self, ConfigError, ConstraintKind, Validate};
31
32use crate::ops::mutation::gaussian_mutation_per_row;
33use crate::rng::{SeedPurpose, seed_stream};
34use crate::strategy::{Strategy, StrategyMetrics};
35
36const DEFAULT_SIGMA_MIN: f32 = 1e-8;
39const DEFAULT_SIGMA_MAX: f32 = 1e6;
42
43#[derive(Debug, Clone)]
45pub struct EpConfig {
46 pub mu: usize,
49 pub genome_dim: usize,
51 pub bounds: Bounds,
53 pub initial_sigma: f32,
55 pub sigma_min: f32,
62 pub sigma_max: f32,
70 pub tau: f32,
73 pub tournament_q: usize,
75}
76
77impl EpConfig {
78 #[must_use]
84 pub fn default_for(mu: usize, genome_dim: usize) -> Self {
85 #[allow(clippy::cast_precision_loss)]
86 let d = genome_dim as f32;
87 let tau = 1.0 / (2.0 * d.sqrt()).sqrt();
88 Self {
89 mu,
90 genome_dim,
91 bounds: Bounds::new(-5.12, 5.12),
92 initial_sigma: 1.0,
93 sigma_min: DEFAULT_SIGMA_MIN,
94 sigma_max: DEFAULT_SIGMA_MAX,
95 tau,
96 tournament_q: 10,
97 }
98 }
99}
100
101impl Validate for EpConfig {
102 fn validate(&self) -> Result<(), ConfigError> {
103 const C: &str = "EpConfig";
104 config::at_least(C, "mu", self.mu, 1)?;
105 config::nonzero(C, "genome_dim", self.genome_dim)?;
106 config::positive(C, "initial_sigma", f64::from(self.initial_sigma))?;
107 config::positive(C, "sigma_min", f64::from(self.sigma_min))?;
108 config::ordered(
109 C,
110 "sigma_max",
111 f64::from(self.sigma_min),
112 f64::from(self.sigma_max),
113 )?;
114 config::positive(C, "tau", f64::from(self.tau))?;
115 config::at_least(C, "tournament_q", self.tournament_q, 1)?;
116 if self.tournament_q > 2 * self.mu {
117 return Err(ConfigError {
118 config: C,
119 field: "tournament_q",
120 kind: ConstraintKind::Custom("tournament_q must not exceed 2 * mu"),
121 });
122 }
123 Ok(())
124 }
125}
126
127#[derive(Debug, Clone)]
141pub struct EpState<B: Backend> {
142 pub parents: Tensor<B, 2>,
144 pub sigmas: Tensor<B, 1>,
147 pub parent_fitness: Vec<f32>,
153 pub best_genome: Option<Tensor<B, 2>>,
157 pub best_fitness: f32,
162 pub generation: usize,
164}
165
166#[derive(Debug, Clone, Copy, Default)]
179pub struct EvolutionaryProgramming<B: Backend> {
180 _backend: PhantomData<fn() -> B>,
181}
182
183impl<B: Backend> EvolutionaryProgramming<B> {
184 #[must_use]
186 pub fn new() -> Self {
187 Self {
188 _backend: PhantomData,
189 }
190 }
191}
192
193impl<B: Backend> Strategy<B> for EvolutionaryProgramming<B>
194where
195 B::Device: Clone,
196{
197 type Params = EpConfig;
198 type State = EpState<B>;
199 type Genome = Tensor<B, 2>;
200
201 fn init(
210 &self,
211 params: &EpConfig,
212 rng: &mut dyn Rng,
213 device: &<B as burn::tensor::backend::BackendTypes>::Device,
214 ) -> EpState<B> {
215 debug_assert!(
216 params.validate().is_ok(),
217 "invalid EpConfig reached init: {params:?}"
218 );
219 let (lo, hi): (f32, f32) = params.bounds.into();
220 let mu = params.mu;
225 let genome_dim = params.genome_dim;
226 let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
227 let mut parent_rows = Vec::with_capacity(mu * genome_dim);
228 for _ in 0..mu * genome_dim {
229 parent_rows.push(lo + (hi - lo) * stream.random::<f32>());
230 }
231 let parents =
232 Tensor::<B, 2>::from_data(TensorData::new(parent_rows, [mu, genome_dim]), device);
233 let sigmas = Tensor::<B, 1>::from_data(
234 TensorData::new(vec![params.initial_sigma; params.mu], [params.mu]),
235 device,
236 );
237 EpState {
238 parents,
239 sigmas,
240 parent_fitness: Vec::new(),
241 best_genome: None,
242 best_fitness: f32::NEG_INFINITY,
243 generation: 0,
244 }
245 }
246
247 fn ask(
267 &self,
268 params: &EpConfig,
269 state: &EpState<B>,
270 rng: &mut dyn Rng,
271 device: &<B as burn::tensor::backend::BackendTypes>::Device,
272 ) -> (Tensor<B, 2>, EpState<B>) {
273 if state.parent_fitness.is_empty() {
275 return (state.parents.clone(), state.clone());
276 }
277
278 let mu = params.mu;
279 let mut sigma_rng =
280 seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);
281 let mut mutation_rng = seed_stream(
282 rng.next_u64(),
283 state.generation as u64,
284 SeedPurpose::Mutation,
285 );
286
287 let mut noise_rows = Vec::with_capacity(mu);
291 for _ in 0..mu {
292 noise_rows.push(crate::sampling::standard_normal(&mut sigma_rng));
293 }
294 let noise = Tensor::<B, 1>::from_data(TensorData::new(noise_rows, [mu]), device);
295 let offspring_sigmas = (state.sigmas.clone() * noise.mul_scalar(params.tau).exp())
299 .clamp(params.sigma_min, params.sigma_max);
300
301 let offspring = gaussian_mutation_per_row(
304 state.parents.clone(),
305 offspring_sigmas.clone(),
306 &mut mutation_rng,
307 device,
308 );
309 let (lo, hi): (f32, f32) = params.bounds.into();
310 let offspring = offspring.clamp(lo, hi);
311
312 let mut state = state.clone();
314 state.sigmas = Tensor::cat(vec![state.sigmas.clone(), offspring_sigmas], 0);
315 (offspring, state)
316 }
317
318 fn tell(
340 &self,
341 params: &EpConfig,
342 offspring: Tensor<B, 2>,
343 fitness: Tensor<B, 1>,
344 mut state: EpState<B>,
345 rng: &mut dyn Rng,
346 ) -> (EpState<B>, StrategyMetrics) {
347 let fitness_host = fitness
348 .into_data()
349 .into_vec::<f32>()
350 .expect("fitness tensor must be readable as f32");
351 let device = offspring.device();
352
353 if state.parent_fitness.is_empty() {
355 state.parent_fitness.clone_from(&fitness_host);
356 state.generation += 1;
357 update_best(&mut state, &offspring, &fitness_host);
358 let m = StrategyMetrics::from_host_fitness(
359 state.generation,
360 &fitness_host,
361 state.best_fitness,
362 );
363 state.best_fitness = m.best_fitness_ever();
364 state.parents = offspring;
365 state.sigmas = Tensor::<B, 1>::from_data(
366 TensorData::new(vec![params.initial_sigma; params.mu], [params.mu]),
367 &device,
368 );
369 return (state, m);
370 }
371
372 let mu = params.mu;
373 let combined_pop = Tensor::cat(vec![state.parents.clone(), offspring.clone()], 0);
375 let combined_fit: Vec<f32> = state
376 .parent_fitness
377 .iter()
378 .chain(fitness_host.iter())
379 .copied()
380 .collect();
381 let combined_sigmas = state.sigmas.clone(); let mut selection_rng = seed_stream(
387 rng.next_u64(),
388 state.generation as u64,
389 SeedPurpose::Selection,
390 );
391 let n = combined_fit.len();
392 let mut win_counts: Vec<u32> = vec![0; n];
393 for (i, &my_fit) in combined_fit.iter().enumerate() {
394 for _ in 0..params.tournament_q {
395 let opp = selection_rng.random_range(0..n);
396 if my_fit > combined_fit[opp] {
397 win_counts[i] += 1;
398 }
399 }
400 }
401
402 let mut indexed: Vec<usize> = (0..n).collect();
405 let sane: Vec<f32> = combined_fit
406 .iter()
407 .map(|&f| crate::fitness::sanitize_fitness(f))
408 .collect();
409 indexed.sort_by(|&a, &b| {
410 win_counts[b]
411 .cmp(&win_counts[a])
412 .then_with(|| sane[b].total_cmp(&sane[a]))
413 });
414 indexed.truncate(mu);
415 #[allow(clippy::cast_possible_wrap)]
416 let survivor_idx: Vec<i64> = indexed.iter().map(|&i| i as i64).collect();
417
418 let idx_tensor =
419 Tensor::<B, 1, Int>::from_data(TensorData::new(survivor_idx.clone(), [mu]), &device);
420 let next_parents = combined_pop.select(0, idx_tensor.clone());
421 let next_sigmas = combined_sigmas.select(0, idx_tensor);
422 let next_fitness: Vec<f32> = survivor_idx
423 .iter()
424 .map(|&i| {
425 #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
426 combined_fit[i as usize]
427 })
428 .collect();
429
430 state.parents = next_parents;
431 state.sigmas = next_sigmas;
432 state.parent_fitness = next_fitness;
433 state.generation += 1;
434 update_best(&mut state, &offspring, &fitness_host);
435 let m =
436 StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
437 state.best_fitness = m.best_fitness_ever();
438 (state, m)
439 }
440
441 fn best(&self, state: &EpState<B>) -> Option<(Tensor<B, 2>, f32)> {
446 state
447 .best_genome
448 .as_ref()
449 .map(|g| (g.clone(), state.best_fitness))
450 }
451}
452
453fn update_best<B: Backend>(state: &mut EpState<B>, pop: &Tensor<B, 2>, fitness: &[f32]) {
454 if fitness.is_empty() {
455 return;
456 }
457 let mut best_idx = 0usize;
458 let mut best_f = fitness[0];
459 for (i, &f) in fitness.iter().enumerate().skip(1) {
460 if f > best_f {
461 best_f = f;
462 best_idx = i;
463 }
464 }
465 if best_f > state.best_fitness {
466 let device = pop.device();
467 #[allow(clippy::cast_possible_wrap)]
468 let idx =
469 Tensor::<B, 1, Int>::from_data(TensorData::new(vec![best_idx as i64], [1]), &device);
470 state.best_genome = Some(pop.clone().select(0, idx));
471 state.best_fitness = best_f;
472 }
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478 use crate::fitness::FromFitnessEvaluable;
479 use crate::strategy::EvolutionaryHarness;
480 use burn::backend::Flex;
481 use rlevo_core::fitness::FitnessEvaluable;
482 type TestBackend = Flex;
483
484 #[test]
485 fn default_config_validates() {
486 assert!(EpConfig::default_for(30, 10).validate().is_ok());
487 }
488
489 #[test]
490 fn rejects_tournament_q_above_two_mu() {
491 let mut cfg = EpConfig::default_for(5, 10);
492 cfg.tournament_q = 11;
493 assert_eq!(cfg.validate().unwrap_err().field, "tournament_q");
494 }
495
496 #[test]
500 fn rejects_zero_mu() {
501 let cfg = EpConfig::default_for(0, 10);
502 assert_eq!(cfg.validate().unwrap_err().field, "mu");
503 }
504
505 #[test]
509 fn mu_one_is_handled() {
510 use rand::SeedableRng;
511 use rand::rngs::StdRng;
512
513 let device = Default::default();
514 let strategy = EvolutionaryProgramming::<TestBackend>::new();
515 let mut params = EpConfig::default_for(1, 3);
516 params.tournament_q = 2;
518 assert!(params.validate().is_ok(), "μ = 1 config must validate");
519
520 let mut rng = StdRng::seed_from_u64(3);
521 let mut state = strategy.init(¶ms, &mut rng, &device);
522 for _ in 0..10 {
523 let (offspring, next) = strategy.ask(¶ms, &state, &mut rng, &device);
524 let fitness = neg_sphere(&offspring);
525 let (advanced, _) = strategy.tell(¶ms, offspring, fitness, next, &mut rng);
526 state = advanced;
527 }
528 assert_eq!(
529 state.parents.dims()[0],
530 1,
531 "μ = 1 must keep a single parent"
532 );
533 }
534
535 fn neg_sphere(pop: &Tensor<TestBackend, 2>) -> Tensor<TestBackend, 1> {
538 let device = pop.device();
539 let [n, d] = pop.dims();
540 let rows: Vec<f32> = pop
541 .clone()
542 .into_data()
543 .into_vec::<f32>()
544 .expect("population host-read of a tensor this test just built");
545 #[allow(clippy::needless_range_loop)]
546 let fit: Vec<f32> = (0..n)
547 .map(|i| -(0..d).map(|j| rows[i * d + j].powi(2)).sum::<f32>())
548 .collect();
549 Tensor::<TestBackend, 1>::from_data(TensorData::new(fit, [n]), &device)
550 }
551
552 fn run_ep_trajectory(seed: u64, gens: usize) -> Vec<f32> {
555 use rand::SeedableRng;
556 use rand::rngs::StdRng;
557
558 let device = Default::default();
559 let strategy = EvolutionaryProgramming::<TestBackend>::new();
560 let params = EpConfig::default_for(8, 3);
561 let mut rng = StdRng::seed_from_u64(seed);
562 let mut state = strategy.init(¶ms, &mut rng, &device);
563 let mut traj = Vec::with_capacity(gens);
564 for _ in 0..gens {
565 let (offspring, next) = strategy.ask(¶ms, &state, &mut rng, &device);
566 let fitness = neg_sphere(&offspring);
567 let (advanced, m) = strategy.tell(¶ms, offspring, fitness, next, &mut rng);
568 traj.push(m.best_fitness_ever());
569 state = advanced;
570 }
571 traj
572 }
573
574 #[test]
579 fn same_seed_reproduces_trajectory() {
580 let a = run_ep_trajectory(2024, 30);
581 let b = run_ep_trajectory(2024, 30);
582 assert_eq!(a, b, "EP trajectory diverged under identical seed");
583 }
584
585 #[test]
590 fn best_genome_matches_best_fitness() {
591 use rand::SeedableRng;
592 use rand::rngs::StdRng;
593
594 let device = Default::default();
595 let strategy = EvolutionaryProgramming::<TestBackend>::new();
596 let params = EpConfig::default_for(6, 3);
597 let mut rng = StdRng::seed_from_u64(5);
598 let state = strategy.init(¶ms, &mut rng, &device);
599 let (parents0, s) = strategy.ask(¶ms, &state, &mut rng, &device);
602 let [n, d] = parents0.dims();
603 #[allow(clippy::cast_precision_loss)]
604 let fit_vec: Vec<f32> = (0..n).map(|i| i as f32).collect();
605 let expected_idx = n - 1;
606 let expected_fit = fit_vec[expected_idx];
607 let parent_rows: Vec<f32> = parents0
608 .clone()
609 .into_data()
610 .into_vec::<f32>()
611 .expect("parent host-read of a tensor this test just built");
612 let expected_genome: Vec<f32> =
613 parent_rows[expected_idx * d..(expected_idx + 1) * d].to_vec();
614
615 let fitness = Tensor::<TestBackend, 1>::from_data(TensorData::new(fit_vec, [n]), &device);
616 let (s, _) = strategy.tell(¶ms, parents0, fitness, s, &mut rng);
617
618 let (genome, best_fit) = strategy.best(&s).expect("best after first tell");
619 approx::assert_relative_eq!(best_fit, expected_fit);
620 let got: Vec<f32> = genome
621 .into_data()
622 .into_vec::<f32>()
623 .expect("best-genome host-read of a tensor this test just built");
624 for (g, e) in got.iter().zip(expected_genome.iter()) {
625 approx::assert_relative_eq!(*g, *e);
626 }
627 }
628
629 struct NanSphere;
631 struct NanSphereFit;
632 impl FitnessEvaluable for NanSphereFit {
633 type Individual = Vec<f64>;
634 type Landscape = NanSphere;
635 fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
636 let s: f64 = x.iter().map(|v| v * v).sum();
637 if x[0] > 0.0 { f64::NAN } else { s }
638 }
639 }
640
641 #[test]
647 fn nan_fitness_never_becomes_best() {
648 let device = Default::default();
649 let params = EpConfig::default_for(20, 4);
650 let fitness_fn = FromFitnessEvaluable::new(NanSphereFit, NanSphere);
651 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
652 EvolutionaryProgramming::<TestBackend>::new(),
653 params,
654 fitness_fn,
655 77,
656 device,
657 40,
658 )
659 .expect("valid params");
660 harness.reset();
661 loop {
662 if harness.step(()).done {
663 break;
664 }
665 }
666 let best = harness.latest_metrics().unwrap().best_fitness_ever();
667 assert!(
668 best.is_finite(),
669 "NaN fitness poisoned best_fitness_ever: {best}"
670 );
671 }
672
673 #[test]
677 fn rejects_zero_genome_dim() {
678 let cfg = EpConfig::default_for(5, 0);
679 assert!(
680 !cfg.tau.is_finite(),
681 "precondition: derived tau is non-finite for genome_dim == 0, got {}",
682 cfg.tau
683 );
684 assert_eq!(
685 cfg.validate().unwrap_err().field,
686 "genome_dim",
687 "genome_dim == 0 must be rejected before the non-finite tau can be used"
688 );
689 }
690
691 #[test]
694 fn rejects_inverted_sigma_window() {
695 let mut cfg = EpConfig::default_for(5, 10);
696 cfg.sigma_min = 10.0;
697 cfg.sigma_max = 1.0;
698 assert_eq!(
699 cfg.validate().unwrap_err().field,
700 "sigma_max",
701 "sigma_min >= sigma_max must be rejected"
702 );
703 }
704
705 #[test]
710 fn sigma_stays_within_bounds_across_updates() {
711 use rand::SeedableRng;
712 use rand::rngs::StdRng;
713
714 let device = Default::default();
715 let strategy = EvolutionaryProgramming::<TestBackend>::new();
716 let mut params = EpConfig::default_for(6, 3);
717 params.tau = 5.0;
720 params.sigma_min = 1e-4;
721 params.sigma_max = 10.0;
722 assert!(params.validate().is_ok(), "test config must be valid");
723
724 let mut rng = StdRng::seed_from_u64(7);
725 let mut state = strategy.init(¶ms, &mut rng, &device);
726 for generation in 0..60 {
727 let (offspring, next) = strategy.ask(¶ms, &state, &mut rng, &device);
728 let sigmas: Vec<f32> = next
729 .sigmas
730 .clone()
731 .into_data()
732 .into_vec::<f32>()
733 .expect("sigma host-read of a tensor this test just built");
734 for &s in &sigmas {
735 assert!(
736 s.is_finite() && s >= params.sigma_min && s <= params.sigma_max,
737 "σ left [{}, {}] at gen {generation}: {s}",
738 params.sigma_min,
739 params.sigma_max
740 );
741 }
742 let n = offspring.dims()[0];
743 let fitness = Tensor::<TestBackend, 1>::from_data(
744 TensorData::new(vec![1.0_f32; n], [n]),
745 &device,
746 );
747 let (advanced, _) = strategy.tell(¶ms, offspring, fitness, next, &mut rng);
748 state = advanced;
749 }
750 }
751
752 struct Sphere;
753 struct SphereFit;
754 impl FitnessEvaluable for SphereFit {
755 type Individual = Vec<f64>;
756 type Landscape = Sphere;
757 fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
758 x.iter().map(|v| v * v).sum()
759 }
760 }
761
762 #[test]
763 fn ep_converges_on_sphere_d2() {
764 let device = Default::default();
765 let params = EpConfig::default_for(10, 2);
766 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
767 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
768 EvolutionaryProgramming::<TestBackend>::new(),
769 params,
770 fitness_fn,
771 3,
772 device,
773 300,
774 )
775 .expect("valid params");
776 harness.reset();
777 loop {
778 if harness.step(()).done {
779 break;
780 }
781 }
782 let best = harness.latest_metrics().unwrap().best_fitness_ever();
783 assert!(best < 1e-2, "EP Sphere-D2 best={best}");
784 }
785
786 #[test]
787 fn ep_converges_on_sphere_d10() {
788 let device = Default::default();
789 let params = EpConfig::default_for(20, 10);
790 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
791 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
792 EvolutionaryProgramming::<TestBackend>::new(),
793 params,
794 fitness_fn,
795 5,
796 device,
797 2000,
798 )
799 .expect("valid params");
800 harness.reset();
801 loop {
802 if harness.step(()).done {
803 break;
804 }
805 }
806 let best = harness.latest_metrics().unwrap().best_fitness_ever();
807 assert!(best < 1e-4, "EP Sphere-D10 best={best}");
808 }
809}