rlevo_evolution/algorithms/metaheuristic/
aco_r.rs

1//! Ant Colony Optimization for continuous domains (`ACO_R`).
2//!
3//! Socha & Dorigo's 2008 extension of ACO to real-valued search spaces.
4//! The colony maintains an **archive** of the `k` best solutions seen
5//! so far; per generation it samples `m` offspring by drawing each
6//! dimension from a Gaussian kernel centred on an archive solution
7//! selected by rank-weighted roulette:
8//!
9//! 1. Compute per-archive, per-dim `σ_{l,d} = ξ · mean_e |x_{e,d} − x_{l,d}|`.
10//! 2. For every offspring `i` and every dimension `d`:
11//!    - Sample an archive index `l ∼ Categorical(w)` where
12//!      `w_l ∝ exp(−(rank_l − 1)² / (2·q²·k²))`.
13//!    - Sample `x_{i,d} ∼ N(x_{l,d}, σ_{l,d})`.
14//! 3. Evaluate offspring, merge with archive, keep top `k`.
15//!
16//! # References
17//!
18//! - Socha & Dorigo (2008), *Ant colony optimization for continuous domains*.
19
20use std::f32::consts::PI;
21use std::marker::PhantomData;
22
23use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
24use rand::Rng;
25use rand::RngExt;
26use rand_distr::{Distribution as RandDistDist, Normal};
27
28use crate::rng::{SeedPurpose, seed_stream};
29use crate::strategy::{Strategy, StrategyMetrics};
30
31/// Static configuration for [`AntColonyReal`].
32#[derive(Debug, Clone)]
33pub struct AcoRConfig {
34    /// Archive size (number of "best" solutions kept). Socha & Dorigo's
35    /// recommended default is `k = 50`.
36    pub archive_size: usize,
37    /// Offspring per generation (`m` in the paper).
38    pub m: usize,
39    /// Genome dimensionality.
40    pub genome_dim: usize,
41    /// Search-space bounds.
42    pub bounds: (f32, f32),
43    /// Exploration scale (`ξ`). Higher → wider sampling. Canonical 0.85.
44    pub xi: f32,
45    /// Rank-weight decay (`q`). Smaller → stronger bias toward top of
46    /// the archive. Canonical `q = 0.01` (sharp) up to `q ≈ 0.5` (flat).
47    pub q: f32,
48}
49
50impl AcoRConfig {
51    /// Default configuration for a given archive size, offspring count, and dimensionality.
52    #[must_use]
53    pub fn default_for(archive_size: usize, m: usize, genome_dim: usize) -> Self {
54        Self {
55            archive_size,
56            m,
57            genome_dim,
58            bounds: (-5.12, 5.12),
59            xi: 0.85,
60            q: 0.1,
61        }
62    }
63
64    /// Steady-state offspring count per generation (`m`). Note that
65    /// the very first generation evaluates the full initial archive
66    /// (`archive_size` rows) instead — only generations ≥ 1 score `m`.
67    #[must_use]
68    pub fn steady_state_pop_size(&self) -> usize {
69        self.m
70    }
71}
72
73/// Generation state for [`AntColonyReal`].
74#[derive(Debug, Clone)]
75pub struct AcoRState<B: Backend> {
76    /// Archive of `k` best solutions, shape `(k, D)`.
77    pub archive: Tensor<B, 2>,
78    /// Host-side archive fitness, sorted ascending (best first) after
79    /// the first `tell`.
80    pub archive_fitness: Vec<f32>,
81    /// Cached archive weights (recomputed only when `q` or `k` change).
82    pub weights: Vec<f32>,
83    /// Best-so-far genome.
84    pub best_genome: Option<Tensor<B, 2>>,
85    /// Best-so-far fitness.
86    pub best_fitness: f32,
87    /// Generation counter.
88    pub generation: usize,
89}
90
91/// Ant Colony Optimization (continuous domains).
92///
93/// # Panics
94///
95/// [`Strategy::init`] panics if `params.archive_size < 2` (the σ
96/// computation needs at least two archive solutions to take a pairwise
97/// distance) or if `params.m < 1` (no offspring to draw).
98///
99/// # Example
100///
101/// ```no_run
102/// use burn::backend::Flex;
103/// use rlevo_evolution::algorithms::metaheuristic::aco_r::{AcoRConfig, AntColonyReal};
104///
105/// let strategy = AntColonyReal::<Flex>::new();
106/// let params = AcoRConfig::default_for(50, 10, 10);
107/// let _ = (strategy, params);
108/// ```
109#[derive(Debug, Clone, Copy, Default)]
110pub struct AntColonyReal<B: Backend> {
111    _backend: PhantomData<fn() -> B>,
112}
113
114impl<B: Backend> AntColonyReal<B> {
115    /// Builds a new (stateless) strategy object.
116    #[must_use]
117    pub fn new() -> Self {
118        Self {
119            _backend: PhantomData,
120        }
121    }
122
123    /// Compute rank-based archive weights `w_l ∝ exp(−(l−1)² / (2·q²·k²))`.
124    fn compute_weights(archive_size: usize, q: f32) -> Vec<f32> {
125        #[allow(clippy::cast_precision_loss)]
126        let k = archive_size as f32;
127        let denom = 2.0 * q * q * k * k;
128        let scale = 1.0 / (q * k * (2.0 * PI).sqrt());
129        let mut w: Vec<f32> = (0..archive_size)
130            .map(|l| {
131                #[allow(clippy::cast_precision_loss)]
132                let rank = l as f32;
133                scale * (-(rank * rank) / denom).exp()
134            })
135            .collect();
136        let total: f32 = w.iter().sum();
137        for v in &mut w {
138            *v /= total;
139        }
140        w
141    }
142}
143
144impl<B: Backend> Strategy<B> for AntColonyReal<B>
145where
146    B::Device: Clone,
147{
148    type Params = AcoRConfig;
149    type State = AcoRState<B>;
150    type Genome = Tensor<B, 2>;
151
152    /// Initialises the archive by host-sampling `archive_size × genome_dim`
153    /// values uniformly from `params.bounds`.
154    ///
155    /// All random draws go through [`seed_stream`] derived from `rng` rather
156    /// than `B::seed` + `Tensor::random`; this keeps draws reproducible across
157    /// thread schedules when multiple tests or harnesses share the same
158    /// process-wide Burn RNG state.
159    ///
160    /// # Panics
161    ///
162    /// Panics if `params.archive_size < 2` or `params.m < 1`; see
163    /// [`AntColonyReal`] struct-level docs for details.
164    fn init(&self, params: &AcoRConfig, rng: &mut dyn Rng, device: &<B as burn::tensor::backend::BackendTypes>::Device) -> AcoRState<B> {
165        assert!(params.archive_size >= 2, "ACO_R requires archive_size >= 2");
166        assert!(params.m >= 1, "ACO_R requires m >= 1");
167        let (lo, hi) = params.bounds;
168        // Host-sample the initial archive from a deterministic `seed_stream`
169        // rather than the process-wide Flex RNG (`B::seed` + `Tensor::random`),
170        // whose draws interleave with sibling tests under the parallel runner
171        // and are not reproducible across thread schedules.
172        let rows = params.archive_size;
173        let genome_dim = params.genome_dim;
174        let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
175        let mut archive_rows = Vec::with_capacity(rows * genome_dim);
176        for _ in 0..rows * genome_dim {
177            archive_rows.push(lo + (hi - lo) * stream.random::<f32>());
178        }
179        let archive =
180            Tensor::<B, 2>::from_data(TensorData::new(archive_rows, [rows, genome_dim]), device);
181        AcoRState {
182            archive,
183            archive_fitness: Vec::new(),
184            weights: Self::compute_weights(params.archive_size, params.q),
185            best_genome: None,
186            best_fitness: f32::INFINITY,
187            generation: 0,
188        }
189    }
190
191    /// Proposes the next population.
192    ///
193    /// **First call** (`state.archive_fitness.is_empty()`): returns the
194    /// initial archive as-is so the harness scores it before any generation
195    /// update occurs.
196    ///
197    /// **Subsequent calls**: draws `m` offspring by, for each dimension `d`
198    /// of each offspring `i`:
199    ///
200    /// 1. Selecting an archive index `l` by CDF-weighted roulette from
201    ///    `state.weights` (host-side, via `seed_stream`).
202    /// 2. Computing `σ_{l,d} = ξ · mean_e |archive_{e,d} − archive_{l,d}|`
203    ///    on-device.
204    /// 3. Sampling `x_{i,d} ~ N(archive_{l,d}, σ_{l,d})` host-side via
205    ///    `rand_distr::Normal` and a second `seed_stream`.
206    ///
207    /// Offspring are clamped to `params.bounds` before upload to the device.
208    /// The returned state is the same object passed in (no mutation occurs in
209    /// `ask`; all state changes are deferred to [`tell`](Self::tell)).
210    #[allow(clippy::many_single_char_names)]
211    fn ask(
212        &self,
213        params: &AcoRConfig,
214        state: &AcoRState<B>,
215        rng: &mut dyn Rng,
216        device: &<B as burn::tensor::backend::BackendTypes>::Device,
217    ) -> (Tensor<B, 2>, AcoRState<B>) {
218        // First call: evaluate the initial archive.
219        if state.archive_fitness.is_empty() {
220            return (state.archive.clone(), state.clone());
221        }
222
223        let k = params.archive_size;
224        let m = params.m;
225        let d = params.genome_dim;
226
227        // σ[l, j] = ξ · (1/(k-1)) · Σ_e |archive[e, j] - archive[l, j]|
228        // Computed on-device by expanding archive along axis 0 to (k, k, d),
229        // taking |a - b|, reducing along axis 0 (the "e" axis).
230        let archive_l = state.archive.clone().unsqueeze_dim::<3>(0); // (1, k, d)
231        let archive_e = state.archive.clone().unsqueeze_dim::<3>(1); // (k, 1, d)
232        let diffs = (archive_l.expand([k, k, d]) - archive_e.expand([k, k, d])).abs();
233        #[allow(clippy::cast_precision_loss)]
234        let inv = params.xi / ((k - 1).max(1) as f32);
235        let sigma = diffs.sum_dim(0).squeeze::<2>().mul_scalar(inv); // (k, d)
236
237        // Weighted index sampling (host-side) — `m · d` independent draws.
238        let mut stream = seed_stream(
239            rng.next_u64(),
240            state.generation as u64,
241            SeedPurpose::Selection,
242        );
243        let mut mean_rows = vec![0f32; m * d];
244        let mut sigma_rows = vec![0f32; m * d];
245
246        // Gather host-side slices for indexing.
247        let archive_host = state.archive.clone().into_data().into_vec::<f32>().unwrap();
248        let sigma_host = sigma.into_data().into_vec::<f32>().unwrap();
249        let cdf: Vec<f32> = {
250            let mut acc = 0.0;
251            let mut v = Vec::with_capacity(k);
252            for &w in &state.weights {
253                acc += w;
254                v.push(acc);
255            }
256            v
257        };
258        let pick = |u: f32| -> usize { cdf.iter().position(|&c| u <= c).unwrap_or(k - 1) };
259
260        for i in 0..m {
261            for j in 0..d {
262                let u: f32 = stream.random::<f32>();
263                let l = pick(u);
264                mean_rows[i * d + j] = archive_host[l * d + j];
265                sigma_rows[i * d + j] = sigma_host[l * d + j].max(1e-12);
266            }
267        }
268
269        // Sample N(mean, sigma) host-side. Using rand_distr keeps the
270        // draw on the same splitmix stream already threaded through
271        // `stream` above.
272        let mut offspring = vec![0f32; m * d];
273        let mut sample_rng = seed_stream(
274            rng.next_u64(),
275            state.generation as u64,
276            SeedPurpose::Mutation,
277        );
278        for (idx, out) in offspring.iter_mut().enumerate() {
279            let normal = Normal::new(mean_rows[idx], sigma_rows[idx]).expect("sigma > 0");
280            *out = normal.sample(&mut sample_rng);
281        }
282        let (lo, hi) = params.bounds;
283        for v in &mut offspring {
284            *v = v.clamp(lo, hi);
285        }
286        let new_pop = Tensor::<B, 2>::from_data(TensorData::new(offspring, [m, d]), device);
287
288        (new_pop, state.clone())
289    }
290
291    /// Merges `population` with the current archive and updates state.
292    ///
293    /// **First tell** (`state.archive_fitness.is_empty()`): `population` is
294    /// the initial archive returned verbatim by the first `ask`; this call
295    /// sorts it by fitness and records `best_genome` and `best_fitness`.
296    ///
297    /// **Steady-state tell**: concatenates `state.archive` (shape `(k, D)`)
298    /// with `population` (shape `(m, D)`), sorts the combined `k + m` rows by
299    /// fitness, and retains only the top `k`. Updates `best_genome` if the new
300    /// archive leader improves on `state.best_fitness`.
301    ///
302    /// Returns the updated state and a [`StrategyMetrics`] snapshot built from
303    /// `fitness` (the offspring scores, not the full archive).
304    fn tell(
305        &self,
306        params: &AcoRConfig,
307        population: Tensor<B, 2>,
308        fitness: Tensor<B, 1>,
309        mut state: AcoRState<B>,
310        _rng: &mut dyn Rng,
311    ) -> (AcoRState<B>, StrategyMetrics) {
312        let fitness_host = fitness.into_data().into_vec::<f32>().unwrap_or_default();
313        let device = population.device();
314        let k = params.archive_size;
315
316        // First tell: the population being scored IS the initial archive.
317        if state.archive_fitness.is_empty() {
318            // Sort archive by fitness.
319            let mut idx: Vec<usize> = (0..fitness_host.len()).collect();
320            idx.sort_by(|&a, &b| fitness_host[a].partial_cmp(&fitness_host[b]).unwrap());
321            #[allow(clippy::cast_possible_wrap)]
322            let sorted_idx = Tensor::<B, 1, Int>::from_data(
323                TensorData::new(idx.iter().map(|&i| i as i64).collect::<Vec<_>>(), [k]),
324                &device,
325            );
326            state.archive = population.clone().select(0, sorted_idx);
327            state.archive_fitness = idx.iter().map(|&i| fitness_host[i]).collect();
328            state.best_fitness = state.archive_fitness[0];
329            let first_idx =
330                Tensor::<B, 1, Int>::from_data(TensorData::new(vec![0_i64], [1]), &device);
331            state.best_genome = Some(state.archive.clone().select(0, first_idx));
332            state.generation += 1;
333            let m = StrategyMetrics::from_host_fitness(
334                state.generation,
335                &fitness_host,
336                state.best_fitness,
337            );
338            state.best_fitness = m.best_fitness_ever;
339            return (state, m);
340        }
341
342        // Steady state: merge archive + offspring, keep top-k.
343        let combined = Tensor::cat(vec![state.archive.clone(), population.clone()], 0);
344        let mut combined_f: Vec<f32> = state.archive_fitness.clone();
345        combined_f.extend_from_slice(&fitness_host);
346        let mut idx: Vec<usize> = (0..combined_f.len()).collect();
347        idx.sort_by(|&a, &b| combined_f[a].partial_cmp(&combined_f[b]).unwrap());
348        idx.truncate(k);
349        #[allow(clippy::cast_possible_wrap)]
350        let top_idx = Tensor::<B, 1, Int>::from_data(
351            TensorData::new(idx.iter().map(|&i| i as i64).collect::<Vec<_>>(), [k]),
352            &device,
353        );
354        state.archive = combined.select(0, top_idx);
355        state.archive_fitness = idx.iter().map(|&i| combined_f[i]).collect();
356
357        if state.archive_fitness[0] < state.best_fitness {
358            state.best_fitness = state.archive_fitness[0];
359            let first_idx =
360                Tensor::<B, 1, Int>::from_data(TensorData::new(vec![0_i64], [1]), &device);
361            state.best_genome = Some(state.archive.clone().select(0, first_idx));
362        }
363
364        state.generation += 1;
365        let m =
366            StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
367        state.best_fitness = m.best_fitness_ever;
368        (state, m)
369    }
370
371    /// Returns the best genome seen across all generations, or `None` before
372    /// the first [`tell`](Self::tell) call.
373    ///
374    /// The returned tensor has shape `(1, genome_dim)`.
375    fn best(&self, state: &AcoRState<B>) -> Option<(Tensor<B, 2>, f32)> {
376        state
377            .best_genome
378            .as_ref()
379            .map(|g| (g.clone(), state.best_fitness))
380    }
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386    use crate::fitness::FromFitnessEvaluable;
387    use crate::strategy::EvolutionaryHarness;
388    use burn::backend::Flex;
389    use rlevo_core::fitness::FitnessEvaluable;
390
391    type TestBackend = Flex;
392
393    struct Sphere;
394    struct SphereFit;
395    impl FitnessEvaluable for SphereFit {
396        type Individual = Vec<f64>;
397        type Landscape = Sphere;
398        fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
399            x.iter().map(|v| v * v).sum()
400        }
401    }
402
403    #[test]
404    fn weights_sum_to_one() {
405        let w = AntColonyReal::<TestBackend>::compute_weights(10, 0.1);
406        let total: f32 = w.iter().sum();
407        approx::assert_relative_eq!(total, 1.0, epsilon = 1e-5);
408    }
409
410    #[test]
411    fn aco_r_converges_on_sphere_d10() {
412        let device = Default::default();
413        let strategy = AntColonyReal::<TestBackend>::new();
414        let params = AcoRConfig::default_for(30, 15, 10);
415        let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
416        let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
417            strategy, params, fitness_fn, 17, device, 400,
418        );
419        harness.reset();
420        while !harness.step(()).done {}
421        let best = harness.latest_metrics().unwrap().best_fitness_ever;
422        assert!(best < 1e-3, "ACO_R D10 best={best}");
423    }
424}
rlevo_evolution/algorithms/metaheuristic/aco_r.rs

rlevo_evolution/algorithms/metaheuristic/
aco_r.rs
