rlevo_evolution/algorithms/neuroevolution/
weight_only.rs

1//! [`WeightOnly`] — a [`Strategy`] wrapper that ties a flat-genome strategy to
2//! a concrete Burn [`Module`] for weight-only neuroevolution.
3//!
4//! `WeightOnly<B, S, M>` composes an inner `S: Strategy<B, Genome =
5//! Tensor<B, 2>>` with a fixed-topology module `M`. It owns a
6//! [`ModuleReshaper`] (the single source of truth for the genome width,
7//! [`num_params`](WeightOnly::num_params)) but performs **no reshaping** in
8//! `ask`/`tell` — the genome stays a flat `(pop_size, num_params)` tensor and
9//! every `Strategy` method delegates verbatim to the inner strategy. Reshaping
10//! to a module happens only at the fitness boundary, in
11//! [`ModuleEvalFn`](crate::module_eval_fn::ModuleEvalFn) (supervised fitness)
12//! or `rlevo_hybrid`'s `RolloutFitness` (RL fitness).
13//!
14//! This follows the [`MemeticWrapper`](crate::algorithms::memetic::MemeticWrapper)
15//! pattern (ADR 0016): the inner strategy owns all population state; the
16//! wrapper adds one orthogonal concern (here, the module binding).
17//!
18//! # Gradient isolation
19//!
20//! `B: Backend`, **not** `AutodiffBackend`. A caller holding an autodiff
21//! module calls `.valid()` before constructing the wrapper, so gradient
22//! tracking cannot leak into evolution — enforced at the type level.
23
24use std::fmt::Debug;
25use std::marker::PhantomData;
26
27use burn::module::Module;
28use burn::tensor::{Tensor, backend::Backend};
29use rand::Rng;
30
31use crate::param_reshaper::ModuleReshaper;
32use crate::strategy::{Strategy, StrategyMetrics};
33
34/// Composes a flat-genome [`Strategy`] with a fixed-topology Burn [`Module`].
35///
36/// # Type Parameters
37///
38/// - `B`: Burn backend (non-autodiff — see module docs).
39/// - `S`: inner strategy with `Genome = Tensor<B, 2>` (GA, ES, DE, …).
40/// - `M`: the network whose weights are evolved.
41///
42/// # Example
43///
44/// ```ignore
45/// let template = MyMlp::<B>::new(&device);
46/// let strategy = WeightOnly::new(GeneticAlgorithm::<B>::new(), template.clone());
47/// let params = GaConfig::default_for(64, strategy.num_params());
48/// // pair with a ModuleEvalFn over the same template in the harness
49/// ```
50pub struct WeightOnly<B, S, M>
51where
52    B: Backend,
53    S: Strategy<B, Genome = Tensor<B, 2>>,
54    M: Module<B>,
55{
56    inner: S,
57    reshaper: ModuleReshaper<B, M>,
58    _backend: PhantomData<fn() -> B>,
59}
60
61impl<B, S, M> WeightOnly<B, S, M>
62where
63    B: Backend,
64    S: Strategy<B, Genome = Tensor<B, 2>>,
65    M: Module<B>,
66{
67    /// Build a wrapper from an inner strategy and a template module.
68    ///
69    /// The template is cloned into a [`ModuleReshaper`]; its float-leaf count
70    /// becomes [`num_params`](Self::num_params), which must equal the inner
71    /// strategy's configured genome width.
72    pub fn new(inner: S, template: M) -> Self {
73        Self {
74            inner,
75            reshaper: ModuleReshaper::new(template),
76            _backend: PhantomData,
77        }
78    }
79
80    /// Genome width — the number of float parameters in the template module.
81    #[must_use]
82    pub fn num_params(&self) -> usize {
83        self.reshaper.num_params()
84    }
85
86    /// Borrow the owned reshaper (e.g. to build a matching fitness adapter).
87    #[must_use]
88    pub fn reshaper(&self) -> &ModuleReshaper<B, M> {
89        &self.reshaper
90    }
91
92    /// Borrow the inner strategy.
93    #[must_use]
94    pub fn inner(&self) -> &S {
95        &self.inner
96    }
97}
98
99impl<B, S, M> Debug for WeightOnly<B, S, M>
100where
101    B: Backend,
102    S: Strategy<B, Genome = Tensor<B, 2>>,
103    M: Module<B>,
104{
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        f.debug_struct("WeightOnly")
107            .field("num_params", &self.reshaper.num_params())
108            .finish_non_exhaustive()
109    }
110}
111
112impl<B, S, M> Strategy<B> for WeightOnly<B, S, M>
113where
114    B: Backend,
115    S: Strategy<B, Genome = Tensor<B, 2>>,
116    // `Sync` is required so the wrapper satisfies the `Strategy: Send + Sync`
117    // supertrait via its `ModuleReshaper<B, M>` field; Burn modules built from
118    // `Param<Tensor>` leaves are `Sync`.
119    M: Module<B> + Sync,
120{
121    type Params = S::Params;
122    type State = S::State;
123    type Genome = Tensor<B, 2>;
124
125    /// Pure delegation to the inner strategy's `init`.
126    fn init(
127        &self,
128        params: &Self::Params,
129        rng: &mut dyn Rng,
130        device: &<B as burn::tensor::backend::BackendTypes>::Device,
131    ) -> Self::State {
132        self.inner.init(params, rng, device)
133    }
134
135    /// Pure delegation to the inner strategy's `ask`. No reshaping here — the
136    /// genome stays a flat `(pop_size, num_params)` tensor.
137    fn ask(
138        &self,
139        params: &Self::Params,
140        state: &Self::State,
141        rng: &mut dyn Rng,
142        device: &<B as burn::tensor::backend::BackendTypes>::Device,
143    ) -> (Self::Genome, Self::State) {
144        self.inner.ask(params, state, rng, device)
145    }
146
147    /// Pure delegation to the inner strategy's `tell`. Reshaping to a module
148    /// already happened in the fitness adapter that produced `fitness`.
149    fn tell(
150        &self,
151        params: &Self::Params,
152        population: Self::Genome,
153        fitness: Tensor<B, 1>,
154        state: Self::State,
155        rng: &mut dyn Rng,
156    ) -> (Self::State, StrategyMetrics) {
157        self.inner.tell(params, population, fitness, state, rng)
158    }
159
160    /// Pure delegation to the inner strategy's `best`.
161    fn best(&self, state: &Self::State) -> Option<(Self::Genome, f32)> {
162        self.inner.best(state)
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::algorithms::de::DifferentialEvolution;
170    use crate::algorithms::es_classical::EvolutionStrategy;
171    use crate::algorithms::ga::{GaConfig, GeneticAlgorithm};
172    use crate::rng::{SeedPurpose, seed_stream};
173    use burn::backend::Flex;
174    use burn::module::Module;
175    use burn::nn::{Linear, LinearConfig};
176
177    type TestBackend = Flex;
178
179    #[derive(Module, Debug)]
180    struct Mlp<B: Backend> {
181        l1: Linear<B>,
182        l2: Linear<B>,
183    }
184
185    impl<B: Backend> Mlp<B> {
186        fn new(device: &B::Device) -> Self {
187            Self {
188                l1: LinearConfig::new(2, 3).init(device),
189                l2: LinearConfig::new(3, 1).init(device),
190            }
191        }
192    }
193
194    /// AC #3: `WeightOnly` composes with GA, ES, and DE inner strategies and
195    /// reports the template's parameter count. `2*3 + 3` + `3*1 + 1` = `13`.
196    /// Confirms the wrapper satisfies the `Strategy<B>` bound for an inner `S`.
197    fn assert_strategy<T: Strategy<TestBackend>>(_: &T) {}
198
199    #[test]
200    fn composes_with_all_phase1_strategies() {
201        let device = Default::default();
202        let template = Mlp::<TestBackend>::new(&device);
203
204        let ga = WeightOnly::new(GeneticAlgorithm::<TestBackend>::new(), template.clone());
205        let es = WeightOnly::new(EvolutionStrategy::<TestBackend>::new(), template.clone());
206        let de = WeightOnly::new(DifferentialEvolution::<TestBackend>::new(), template);
207
208        assert_eq!(ga.num_params(), 13);
209        assert_eq!(es.num_params(), 13);
210        assert_eq!(de.num_params(), 13);
211
212        assert_strategy(&ga);
213        assert_strategy(&es);
214        assert_strategy(&de);
215    }
216
217    /// §7.1: drives the full `Strategy` surface (`init → ask → tell → best`)
218    /// through the wrapper, proving every method delegates verbatim rather than
219    /// only that `num_params` is reported. The existing test above checks the
220    /// width; this one checks the *plumbing*.
221    #[test]
222    fn test_weight_only_delegates_init_ask_tell_roundtrip() {
223        let device = Default::default();
224        let template = Mlp::<TestBackend>::new(&device);
225        let strategy = WeightOnly::new(GeneticAlgorithm::<TestBackend>::new(), template.clone());
226        let num_params: usize = strategy.num_params();
227
228        let params: GaConfig = GaConfig::default_for(8, num_params);
229        let mut rng = seed_stream(0, 0, SeedPurpose::Init);
230
231        let state = strategy.init(&params, &mut rng, &device);
232        let (pop, state) = strategy.ask(&params, &state, &mut rng, &device);
233        assert_eq!(
234            pop.dims(),
235            [8, num_params],
236            "ask must return a (pop_size, num_params) population"
237        );
238
239        let fitness = Tensor::<TestBackend, 1>::full([8], 1.0, &device);
240        let (state, metrics) = strategy.tell(&params, pop, fitness, state, &mut rng);
241        assert_eq!(
242            metrics.population_size(),
243            8,
244            "tell metrics must report the delegated population size"
245        );
246
247        // `best` is pure delegation; after one `tell` the inner GA has recorded
248        // an elite genome.
249        assert!(
250            strategy.best(&state).is_some(),
251            "best must exist after one tell"
252        );
253
254        // New capability (Change 1): the owned reshaper is `Clone`, and the
255        // clone agrees on the genome width by construction — the single-source
256        // guarantee the wrapper exists to provide.
257        assert_eq!(
258            strategy.reshaper().clone().num_params(),
259            strategy.num_params(),
260            "cloned reshaper must report the same genome width"
261        );
262    }
263}
rlevo_evolution/algorithms/neuroevolution/weight_only.rs

rlevo_evolution/algorithms/neuroevolution/
weight_only.rs
