use std::fmt::Debug;
use std::marker::PhantomData;
use burn::tensor::backend::Backend;
use rand::SeedableRng;
use rand::rngs::StdRng;
use rlevo_core::config::{ConfigError, Validate};
use rlevo_core::evaluation::{BenchEnv, BenchError, BenchStep};
use super::CoEvolutionaryAlgorithm;
#[derive(Debug, Clone)]
pub struct CoEAMetrics {
pub generation: u64,
pub best_fitness_a: f32,
pub best_fitness_b: f32,
pub mean_fitness_a: f32,
pub mean_fitness_b: f32,
pub binding_fitness: f32,
pub hof_size_a: usize,
pub hof_size_b: usize,
}
pub struct CoEvolutionaryHarness<B, C>
where
B: Backend,
C: CoEvolutionaryAlgorithm<B>,
{
algorithm: C,
params: C::Params,
state: Option<C::State>,
rng: StdRng,
base_seed: u64,
device: B::Device,
generation: usize,
max_generations: usize,
latest_metrics: Option<CoEAMetrics>,
_backend: PhantomData<B>,
}
impl<B, C> Debug for CoEvolutionaryHarness<B, C>
where
B: Backend,
C: CoEvolutionaryAlgorithm<B>,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CoEvolutionaryHarness")
.field("base_seed", &self.base_seed)
.field("generation", &self.generation)
.field("max_generations", &self.max_generations)
.field("latest_metrics", &self.latest_metrics)
.finish_non_exhaustive()
}
}
impl<B, C> CoEvolutionaryHarness<B, C>
where
B: Backend,
C: CoEvolutionaryAlgorithm<B>,
{
pub fn new(
algorithm: C,
params: C::Params,
seed: u64,
device: B::Device,
max_generations: usize,
) -> Result<Self, ConfigError>
where
C::Params: Validate,
{
params.validate()?;
Ok(Self {
algorithm,
params,
state: None,
rng: StdRng::seed_from_u64(seed),
base_seed: seed,
device,
generation: 0,
max_generations,
latest_metrics: None,
_backend: PhantomData,
})
}
#[must_use]
pub fn latest_metrics(&self) -> Option<&CoEAMetrics> {
self.latest_metrics.as_ref()
}
#[must_use]
pub fn generation(&self) -> usize {
self.generation
}
pub fn reset(&mut self) {
self.rng = StdRng::seed_from_u64(self.base_seed);
self.generation = 0;
self.latest_metrics = None;
self.state = Some(
self.algorithm
.init(&self.params, &mut self.rng, &self.device),
);
}
pub fn step(&mut self, _action: ()) -> BenchStep<()> {
let state = self
.state
.take()
.expect("CoEvolutionaryHarness::reset must be called before step");
let (new_state, metrics) =
self.algorithm
.step(&self.params, state, &mut self.rng, &self.device);
self.state = Some(new_state);
self.generation += 1;
let reward = f64::from(metrics.binding_fitness);
tracing::info!(
generation = metrics.generation,
best_fitness_a = f64::from(metrics.best_fitness_a),
best_fitness_b = f64::from(metrics.best_fitness_b),
mean_fitness_a = f64::from(metrics.mean_fitness_a),
mean_fitness_b = f64::from(metrics.mean_fitness_b),
hof_size_a = metrics.hof_size_a,
hof_size_b = metrics.hof_size_b,
"coevolution generation",
);
self.latest_metrics = Some(metrics);
let done = self.generation >= self.max_generations;
BenchStep {
observation: (),
reward,
done,
}
}
}
impl<B, C> BenchEnv for CoEvolutionaryHarness<B, C>
where
B: Backend,
C: CoEvolutionaryAlgorithm<B>,
{
type Observation = ();
type Action = ();
fn reset(&mut self) -> Result<Self::Observation, BenchError> {
CoEvolutionaryHarness::<B, C>::reset(self);
Ok(())
}
fn step(&mut self, action: Self::Action) -> Result<BenchStep<Self::Observation>, BenchError> {
Ok(CoEvolutionaryHarness::<B, C>::step(self, action))
}
}
#[cfg(test)]
mod tests {
use super::*;
use burn::backend::Flex;
use burn::tensor::{Tensor, TensorData};
use rlevo_core::bounds::Bounds;
use rlevo_core::objective::ObjectiveSense;
use rlevo_core::probability::Probability;
use rlevo_core::rate::NonNegativeRate;
use crate::algorithms::ga::{
GaConfig, GaCrossover, GaReplacement, GaSelection, GeneticAlgorithm,
};
use crate::coevolution::{CompetitiveCoEA, CompetitiveCoEAParams, CoupledFitness};
type TB = Flex;
const POP: usize = 4;
const DIM: usize = 2;
fn ga_config() -> GaConfig {
GaConfig {
pop_size: POP,
genome_dim: DIM,
bounds: Bounds::new(0.0, 1.0),
mutation_sigma: NonNegativeRate::new(0.1),
selection: GaSelection::Tournament { size: 2 },
crossover: GaCrossover::Uniform {
p: Probability::new(0.5),
},
replacement: GaReplacement::Elitist { elitism_k: 1 },
}
}
struct PoisonRow0Nan;
impl CoupledFitness<TB> for PoisonRow0Nan {
fn evaluate_coupled(&self, populations: &[Tensor<TB, 2>]) -> Vec<Tensor<TB, 1>> {
populations
.iter()
.map(|p| {
let n = p.dims()[0];
let device = p.device();
#[allow(clippy::cast_precision_loss)]
let v: Vec<f32> = (0..n)
.map(|i| if i == 0 { f32::NAN } else { i as f32 })
.collect();
Tensor::<TB, 1>::from_data(TensorData::new(v, [n]), &device)
})
.collect()
}
fn sense(&self) -> ObjectiveSense {
ObjectiveSense::Maximize
}
}
#[test]
fn harness_reward_is_never_nan_with_nan_fitness() {
let device = Default::default();
let algo = CompetitiveCoEA::new(
GeneticAlgorithm::<TB>::new(),
GeneticAlgorithm::<TB>::new(),
PoisonRow0Nan,
);
let params: CompetitiveCoEAParams<GaConfig, GaConfig> = CompetitiveCoEAParams {
params_a: ga_config(),
params_b: ga_config(),
};
let mut harness =
CoEvolutionaryHarness::<TB, _>::new(algo, params, 7, device, 3).expect("valid params");
harness.reset();
let step = harness.step(());
assert!(!step.reward.is_nan(), "harness reward must never be NaN");
assert!(
step.reward.is_finite(),
"reward should be the finite binding value, got {}",
step.reward
);
#[allow(clippy::cast_precision_loss)]
let expected = f64::from((POP - 1) as f32);
approx::assert_relative_eq!(step.reward, expected, epsilon = 1e-6);
}
struct RowCostMin;
impl CoupledFitness<TB> for RowCostMin {
fn evaluate_coupled(&self, populations: &[Tensor<TB, 2>]) -> Vec<Tensor<TB, 1>> {
populations
.iter()
.map(|p| {
let n = p.dims()[0];
let device = p.device();
#[allow(clippy::cast_precision_loss)]
let v: Vec<f32> = (0..n).map(|i| i as f32 + 1.0).collect();
Tensor::<TB, 1>::from_data(TensorData::new(v, [n]), &device)
})
.collect()
}
fn sense(&self) -> ObjectiveSense {
ObjectiveSense::Minimize
}
}
#[test]
fn minimize_harness_reward_is_canonical_binding() {
let device = Default::default();
let algo = CompetitiveCoEA::new(
GeneticAlgorithm::<TB>::new(),
GeneticAlgorithm::<TB>::new(),
RowCostMin,
);
let params: CompetitiveCoEAParams<GaConfig, GaConfig> = CompetitiveCoEAParams {
params_a: ga_config(),
params_b: ga_config(),
};
let mut harness =
CoEvolutionaryHarness::<TB, _>::new(algo, params, 7, device, 3).expect("valid params");
harness.reset();
let step = harness.step(());
assert!(step.reward.is_finite(), "reward must be finite");
approx::assert_relative_eq!(step.reward, -1.0, epsilon = 1e-6);
let m = harness.latest_metrics().expect("metrics after a step");
approx::assert_relative_eq!(m.binding_fitness, -1.0, epsilon = 1e-6);
approx::assert_relative_eq!(m.best_fitness_a, 1.0, epsilon = 1e-6);
}
}