rlevo-evolution 0.2.0

//! Classical Evolution Strategies.
//!
//! Four canonical variants parameterized on a single [`EsConfig`]:
//!
//! - `(1+1)` — a single parent, a single offspring, 1/5th success-rule
//!   σ adaptation.
//! - `(1+λ)` — a single parent, λ offspring per generation; the best
//!   offspring replaces the parent iff its fitness improves. The
//!   underlying mutation/selection loop is also reused by Cartesian GP.
//! - `(μ,λ)` — μ parents, λ offspring; parents are discarded each
//!   generation.
//! - `(μ+λ)` — μ parents, λ offspring; survivors are the μ best of the
//!   combined pool.
//!
//! σ adaptation is by log-normal self-adaptation in the multi-parent
//! variants; `(1+1)` uses Rechenberg's 1/5th success rule.
//!
//! # References
//!
//! - Beyer & Schwefel (2002), *Evolution strategies: A comprehensive
//!   introduction*.

use std::marker::PhantomData;

use burn::tensor::{Tensor, TensorData, backend::Backend};
use rand::Rng;
use rand::RngExt;
use rand_distr::{Distribution as _, Normal};

use crate::ops::mutation::gaussian_mutation_per_row;
use crate::ops::replacement::{mu_comma_lambda, mu_plus_lambda};
use crate::rng::{SeedPurpose, seed_stream};
use crate::strategy::{Strategy, StrategyMetrics};

/// Which selection scheme the ES uses.
#[derive(Debug, Clone, Copy)]
pub enum EsKind {
    /// `(1+1)` with 1/5-rule σ adaptation.
    OnePlusOne,
    /// `(1+λ)` with shared σ across offspring.
    OnePlusLambda { lambda: usize },
    /// `(μ,λ)` with log-normal per-individual σ adaptation.
    MuCommaLambda { mu: usize, lambda: usize },
    /// `(μ+λ)` with log-normal per-individual σ adaptation.
    MuPlusLambda { mu: usize, lambda: usize },
}

impl EsKind {
    /// Returns the effective offspring-population size for this variant.
    #[must_use]
    pub fn population_size(&self) -> usize {
        match self {
            EsKind::OnePlusOne => 1,
            EsKind::OnePlusLambda { lambda }
            | EsKind::MuCommaLambda { lambda, .. }
            | EsKind::MuPlusLambda { lambda, .. } => *lambda,
        }
    }
}

/// Static configuration for an [`EvolutionStrategy`] run.
#[derive(Debug, Clone)]
pub struct EsConfig {
    /// Variant to run.
    pub kind: EsKind,
    /// Genome dimensionality.
    pub genome_dim: usize,
    /// Search-space bounds; used for initialization and clamping.
    pub bounds: (f32, f32),
    /// Initial σ (log-normal self-adaptation modifies it in state).
    pub initial_sigma: f32,
    /// Learning-rate scale for log-normal σ update. Standard default is
    /// `1.0 / sqrt(2 * sqrt(D))`.
    pub tau: f32,
}

impl EsConfig {
    /// Default configuration for a given ES variant and dimensionality.
    ///
    /// Sets `bounds = (-5.12, 5.12)` (the standard Rastrigin/sphere domain),
    /// `initial_sigma = 1.0`, and τ via the standard formula
    /// `1 / sqrt(2 · sqrt(D))` (Beyer & Schwefel 2002, eq. 12).
    #[must_use]
    pub fn default_for(kind: EsKind, genome_dim: usize) -> Self {
        #[allow(clippy::cast_precision_loss)]
        let d = genome_dim as f32;
        let tau = 1.0 / (2.0 * d.sqrt()).sqrt();
        Self {
            kind,
            genome_dim,
            bounds: (-5.12, 5.12),
            initial_sigma: 1.0,
            tau,
        }
    }
}

/// Generation state for [`EvolutionStrategy`].
#[derive(Debug, Clone)]
pub struct EsState<B: Backend> {
    /// Parent population. `(μ, D)` for μ-parent variants; `(1, D)` for
    /// (1+1) and (1+λ).
    pub parents: Tensor<B, 2>,
    /// Per-parent σ values.
    ///
    /// Shape between generations is `(μ,)` for log-normal adaptation and
    /// `(1,)` for `(1+1)`/`(1+λ)` with shared σ. Between an `ask` and the
    /// matching `tell` the tensor is temporarily `(μ + λ,)`: parent σ
    /// followed by per-offspring σ. See `ask` for the rationale.
    pub sigmas: Tensor<B, 1>,
    /// Parent fitnesses.
    pub parent_fitness: Vec<f32>,
    /// Best-so-far genome, shape `(1, D)`.
    pub best_genome: Option<Tensor<B, 2>>,
    /// Best-so-far fitness.
    pub best_fitness: f32,
    /// Completed-generation counter.
    pub generation: usize,
    /// (1+1) only: running success-rate counter for the 1/5th rule.
    pub successes_in_window: u32,
    /// (1+1) only: window length observed so far.
    pub window_len: u32,
}

/// Classical Evolution Strategy.
///
/// # Example
///
/// ```no_run
/// use burn::backend::Flex;
/// use rlevo_evolution::algorithms::es_classical::{EsConfig, EsKind, EvolutionStrategy};
///
/// let strategy = EvolutionStrategy::<Flex>::new();
/// let params = EsConfig::default_for(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10);
/// let _ = (strategy, params);
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct EvolutionStrategy<B: Backend> {
    _backend: PhantomData<fn() -> B>,
}

impl<B: Backend> EvolutionStrategy<B> {
    /// Builds a new (stateless) strategy object.
    #[must_use]
    pub fn new() -> Self {
        Self {
            _backend: PhantomData,
        }
    }

    fn mu(kind: EsKind) -> usize {
        match kind {
            EsKind::OnePlusOne | EsKind::OnePlusLambda { .. } => 1,
            EsKind::MuCommaLambda { mu, .. } | EsKind::MuPlusLambda { mu, .. } => mu,
        }
    }

    fn sample_initial_parents(
        params: &EsConfig,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (Tensor<B, 2>, Tensor<B, 1>) {
        let mu = Self::mu(params.kind);
        let (lo, hi) = params.bounds;
        // Host-sample the initial parents from a deterministic `seed_stream`
        // rather than the process-wide Flex RNG (`B::seed` + `Tensor::random`),
        // whose draws interleave with sibling tests under the parallel runner
        // and are not reproducible across thread schedules.
        let genome_dim = params.genome_dim;
        let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
        let mut parent_rows = Vec::with_capacity(mu * genome_dim);
        for _ in 0..mu * genome_dim {
            parent_rows.push(lo + (hi - lo) * stream.random::<f32>());
        }
        let parents =
            Tensor::<B, 2>::from_data(TensorData::new(parent_rows, [mu, genome_dim]), device);
        let sigmas = Tensor::<B, 1>::from_data(
            TensorData::new(vec![params.initial_sigma; mu], [mu]),
            device,
        );
        (parents, sigmas)
    }
}

impl<B: Backend> Strategy<B> for EvolutionStrategy<B>
where
    B::Device: Clone,
{
    type Params = EsConfig;
    type State = EsState<B>;
    type Genome = Tensor<B, 2>;

    /// Samples the initial parent population uniformly from `params.bounds`
    /// via a deterministic `seed_stream` (host-RNG convention) and
    /// initializes all σ values to `params.initial_sigma`.
    fn init(&self, params: &EsConfig, rng: &mut dyn Rng, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> EsState<B> {
        let (parents, sigmas) = Self::sample_initial_parents(params, rng, device);
        EsState {
            parents,
            sigmas,
            parent_fitness: Vec::new(),
            best_genome: None,
            best_fitness: f32::INFINITY,
            generation: 0,
            successes_in_window: 0,
            window_len: 0,
        }
    }

    /// Generates the offspring population for the current generation.
    ///
    /// On the very first call (before any `tell`), returns the initial parents
    /// unchanged so that they can be fitness-evaluated as the seed population.
    /// On subsequent calls, duplicates parents by uniform random selection,
    /// applies log-normal σ adaptation (multi-parent variants) or inherits the
    /// shared σ (`(1+1)` / `(1+λ)`), then mutates via per-individual Gaussian
    /// noise. All stochastic draws go through `seed_stream`
    /// (host-RNG convention); offspring σ values are appended to
    /// `state.sigmas` for consumption by `tell`.
    fn ask(
        &self,
        params: &EsConfig,
        state: &EsState<B>,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (Tensor<B, 2>, EsState<B>) {
        // First call: evaluate the initial parents as the "offspring"
        // so fitness is populated in the subsequent `tell`.
        if state.parent_fitness.is_empty() {
            return (state.parents.clone(), state.clone());
        }

        let lambda = params.kind.population_size();
        let mu = Self::mu(params.kind);

        let mut mutation_rng = seed_stream(
            rng.next_u64(),
            state.generation as u64,
            SeedPurpose::Mutation,
        );
        let mut sigma_rng =
            seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);

        // Build an offspring population of size λ by sampling a parent
        // index per offspring and mutating. Uniform random parent
        // selection — no fitness pressure applied at this stage in
        // classical ES; survivor selection provides the pressure.
        let mut parent_indices: Vec<i64> = Vec::with_capacity(lambda);
        for _ in 0..lambda {
            #[allow(clippy::cast_possible_wrap)]
            parent_indices.push(sigma_rng.random_range(0..mu) as i64);
        }
        let idx_tensor = Tensor::<B, 1, burn::tensor::Int>::from_data(
            TensorData::new(parent_indices.clone(), [lambda]),
            device,
        );
        let duplicated_parents = state.parents.clone().select(0, idx_tensor.clone());
        let duplicated_sigmas = state.sigmas.clone().select(0, idx_tensor);

        // Apply log-normal σ adaptation (multi-parent case) or keep σ
        // shared (1+1 / 1+λ). Log-normal: σ' = σ * exp(τ · N(0,1)).
        let is_one_plus = matches!(
            params.kind,
            EsKind::OnePlusOne | EsKind::OnePlusLambda { .. }
        );
        let offspring_sigmas = if is_one_plus {
            duplicated_sigmas
        } else {
            // Host-sample the N(0,1) noise from the deterministic `sigma_rng`
            // so the log-normal σ update is reproducible across schedules.
            let normal = Normal::new(0.0f32, 1.0).expect("unit normal is well-defined");
            let mut noise_rows = Vec::with_capacity(lambda);
            for _ in 0..lambda {
                noise_rows.push(normal.sample(&mut sigma_rng));
            }
            let noise = Tensor::<B, 1>::from_data(TensorData::new(noise_rows, [lambda]), device);
            duplicated_sigmas * noise.mul_scalar(params.tau).exp()
        };

        // Mutate parents by the per-offspring σ, drawing from the host
        // `mutation_rng`.
        let mutated = gaussian_mutation_per_row(
            duplicated_parents,
            offspring_sigmas.clone(),
            &mut mutation_rng,
            device,
        );

        // Clamp to bounds.
        let (lo, hi) = params.bounds;
        let mutated = mutated.clamp(lo, hi);

        let mut state = state.clone();
        // Carry offspring σ to `tell` by appending them to `state.sigmas`.
        // After this point sigmas has shape `(μ + λ,)`: the first μ entries
        // are the unchanged parent σ, the last λ are the per-offspring σ.
        // `tell` slices both halves to align survivor σ with survivor genomes
        // (`(μ+λ)` selection draws from the union, `(μ,λ)` only from the λ
        // offspring slice). Folding the offspring σ into the existing field
        // avoids adding a transient pending-σ field to `EsState`.
        let combined_sigmas = Tensor::cat(vec![state.sigmas.clone(), offspring_sigmas], 0);
        state.sigmas = combined_sigmas;
        (mutated, state)
    }

    /// Applies variant-specific selection and σ adaptation, then returns the
    /// updated state and a per-generation metrics snapshot.
    ///
    /// Variant behaviour:
    /// - `(1+1)`: greedy replacement; σ updated by Rechenberg's 1/5th
    ///   success rule every `10·D` steps.
    /// - `(1+λ)`: best offspring replaces the parent only if it strictly
    ///   improves fitness; σ is carried over unchanged.
    /// - `(μ,λ)`: selects the μ best offspring; parent pool discarded.
    ///   Survivor σ values are gathered by the same truncation indices.
    /// - `(μ+λ)`: selects the μ best of the combined parent + offspring
    ///   pool. Survivor σ values are drawn from the concatenated σ vector
    ///   by the same indices.
    ///
    /// The first `tell` after `init` bootstraps `parent_fitness` from the
    /// initial-population evaluation rather than running selection.
    #[allow(clippy::too_many_lines)]
    fn tell(
        &self,
        params: &EsConfig,
        offspring: Tensor<B, 2>,
        fitness: Tensor<B, 1>,
        mut state: EsState<B>,
        _rng: &mut dyn Rng,
    ) -> (EsState<B>, StrategyMetrics) {
        let fitness_host = fitness.into_data().into_vec::<f32>().unwrap_or_default();

        // First `tell` after `init`: offspring here is actually the
        // initial parent population evaluated.
        if state.parent_fitness.is_empty() {
            state.parent_fitness.clone_from(&fitness_host);
            state.generation += 1;
            update_best(&mut state, &offspring, &fitness_host);
            let m = StrategyMetrics::from_host_fitness(
                state.generation,
                &fitness_host,
                state.best_fitness,
            );
            state.best_fitness = m.best_fitness_ever;
            state.parents = offspring;
            // Restore parent-count σ vector.
            let mu = Self::mu(params.kind);
            let device = state.parents.device();
            state.sigmas = Tensor::<B, 1>::from_data(
                TensorData::new(vec![params.initial_sigma; mu], [mu]),
                &device,
            );
            return (state, m);
        }

        let device = offspring.device();
        let mu = Self::mu(params.kind);
        // state.sigmas currently holds parent σ concatenated with
        // offspring σ, per `ask`'s scratchpad trick.
        let lambda = params.kind.population_size();
        #[allow(clippy::single_range_in_vec_init)]
        let parent_sigmas = state.sigmas.clone().slice([0..mu]);
        #[allow(clippy::single_range_in_vec_init)]
        let offspring_sigmas = state.sigmas.clone().slice([mu..(mu + lambda)]);

        match params.kind {
            EsKind::OnePlusOne => {
                // One parent, one offspring. Fitness[0] is the offspring.
                let parent_fit = state.parent_fitness[0];
                let offspring_fit = fitness_host[0];
                let success = offspring_fit < parent_fit;
                state.window_len += 1;
                if success {
                    state.successes_in_window += 1;
                    state.parents.clone_from(&offspring);
                    state.parent_fitness = vec![offspring_fit];
                }
                // Rechenberg 1/5-rule every 10 · D generations.
                #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
                let window = 10_u32.saturating_mul(params.genome_dim as u32).max(1);
                if state.window_len >= window {
                    #[allow(clippy::cast_precision_loss)]
                    let rate = state.successes_in_window as f32 / state.window_len as f32;
                    let current_sigma =
                        state.sigmas.clone().into_data().into_vec::<f32>().unwrap()[0];
                    let new_sigma = if rate > 0.2 {
                        current_sigma * 1.22
                    } else if rate < 0.2 {
                        current_sigma / 1.22
                    } else {
                        current_sigma
                    };
                    state.sigmas =
                        Tensor::<B, 1>::from_data(TensorData::new(vec![new_sigma], [1]), &device);
                    state.successes_in_window = 0;
                    state.window_len = 0;
                } else {
                    state.sigmas = parent_sigmas;
                }
            }
            EsKind::OnePlusLambda { .. } => {
                // Best of (parent, offspring pool).
                let best_off_idx = argmin(&fitness_host);
                let best_off_fit = fitness_host[best_off_idx];
                if best_off_fit < state.parent_fitness[0] {
                    #[allow(clippy::single_range_in_vec_init)]
                    let best_row = offspring.clone().slice([best_off_idx..best_off_idx + 1]);
                    state.parents = best_row;
                    state.parent_fitness = vec![best_off_fit];
                }
                state.sigmas = parent_sigmas;
            }
            EsKind::MuCommaLambda { mu, .. } => {
                let (survivors, survivor_f) =
                    mu_comma_lambda::<B>(offspring.clone(), &fitness_host, mu, &device);
                // Gather survivor σs matching the same indices.
                let survivor_idx =
                    crate::ops::selection::truncation_indices_host(&fitness_host, mu);
                let survivor_sigmas = offspring_sigmas.select(
                    0,
                    Tensor::<B, 1, burn::tensor::Int>::from_data(
                        TensorData::new(survivor_idx, [mu]),
                        &device,
                    ),
                );
                state.parents = survivors;
                state.parent_fitness = survivor_f;
                state.sigmas = survivor_sigmas;
            }
            EsKind::MuPlusLambda { mu, .. } => {
                let (survivors, survivor_f) = mu_plus_lambda::<B>(
                    state.parents.clone(),
                    &state.parent_fitness,
                    offspring.clone(),
                    &fitness_host,
                    mu,
                    &device,
                );
                // Survivor σ via truncation_indices_host on the combined fitness.
                let combined_f: Vec<f32> = state
                    .parent_fitness
                    .iter()
                    .chain(fitness_host.iter())
                    .copied()
                    .collect();
                let survivor_idx = crate::ops::selection::truncation_indices_host(&combined_f, mu);
                let combined_sigmas = Tensor::cat(vec![parent_sigmas, offspring_sigmas], 0);
                let survivor_sigmas = combined_sigmas.select(
                    0,
                    Tensor::<B, 1, burn::tensor::Int>::from_data(
                        TensorData::new(survivor_idx, [mu]),
                        &device,
                    ),
                );
                state.parents = survivors;
                state.parent_fitness = survivor_f;
                state.sigmas = survivor_sigmas;
            }
        }

        state.generation += 1;
        update_best(&mut state, &offspring, &fitness_host);
        let m =
            StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
        state.best_fitness = m.best_fitness_ever;
        (state, m)
    }

    /// Returns the best-so-far genome and its fitness, or `None` before the
    /// first `tell` call.
    fn best(&self, state: &EsState<B>) -> Option<(Tensor<B, 2>, f32)> {
        state
            .best_genome
            .as_ref()
            .map(|g| (g.clone(), state.best_fitness))
    }
}

fn argmin(xs: &[f32]) -> usize {
    let mut best_idx = 0usize;
    let mut best = f32::INFINITY;
    for (i, &v) in xs.iter().enumerate() {
        if v < best {
            best = v;
            best_idx = i;
        }
    }
    best_idx
}

fn update_best<B: Backend>(state: &mut EsState<B>, pop: &Tensor<B, 2>, fitness: &[f32]) {
    if fitness.is_empty() {
        return;
    }
    let best_idx = argmin(fitness);
    let best_f = fitness[best_idx];
    if best_f < state.best_fitness {
        let device = pop.device();
        #[allow(clippy::cast_possible_wrap)]
        let idx = Tensor::<B, 1, burn::tensor::Int>::from_data(
            TensorData::new(vec![best_idx as i64], [1]),
            &device,
        );
        state.best_genome = Some(pop.clone().select(0, idx));
        state.best_fitness = best_f;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fitness::FromFitnessEvaluable;
    use crate::strategy::EvolutionaryHarness;
    use burn::backend::Flex;
    use rlevo_core::fitness::FitnessEvaluable;
    type TestBackend = Flex;

    struct Sphere;
    struct SphereFit;
    impl FitnessEvaluable for SphereFit {
        type Individual = Vec<f64>;
        type Landscape = Sphere;
        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
            x.iter().map(|v| v * v).sum()
        }
    }

    fn run_es(kind: EsKind, dim: usize, generations: usize, seed: u64) -> f32 {
        let device = Default::default();
        let strategy = EvolutionStrategy::<TestBackend>::new();
        let params = EsConfig::default_for(kind, dim);
        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
            strategy,
            params,
            fitness_fn,
            seed,
            device,
            generations,
        );
        harness.reset();
        loop {
            let step = harness.step(());
            if step.done {
                break;
            }
        }
        harness.latest_metrics().unwrap().best_fitness_ever
    }

    #[test]
    fn one_plus_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::OnePlusLambda { lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-2, "OnePlusLambda best={best}");
    }

    #[test]
    fn one_plus_one_converges_on_sphere_d2() {
        let best = run_es(EsKind::OnePlusOne, 2, 500, 11);
        assert!(best < 1e-2, "OnePlusOne best={best}");
    }

    #[test]
    fn mu_plus_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::MuPlusLambda { mu: 3, lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-2, "MuPlusLambda best={best}");
    }

    #[test]
    fn mu_comma_lambda_converges_on_sphere_d2() {
        let best = run_es(EsKind::MuCommaLambda { mu: 3, lambda: 8 }, 2, 200, 7);
        assert!(best < 1e-1, "MuCommaLambda best={best}");
    }

    #[test]
    fn mu_plus_lambda_converges_on_sphere_d10() {
        // Convergence on Sphere (D=10) to best_fitness < 1e-6 within
        // budget on Flex. We allow a generous budget because the
        // classical ES is slower than CMA-ES; the goal is to verify
        // convergence direction, not to optimize hyperparameters.
        let best = run_es(EsKind::MuPlusLambda { mu: 5, lambda: 20 }, 10, 1500, 42);
        assert!(best < 1e-6, "MuPlusLambda D10 best={best}");
    }
}