plasticity_lab/
observer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Per-step session observer types for [`crate::PlasticityTrainer::run_session_with_observer`].
4//!
5//! These types are borrowed snapshots. They never expose `&mut` access to the
6//! network, and constructing a [`TrainingStepEvent`] does not heap-allocate.
7
8use neuromod::NeuroModulators;
9
10/// Borrowed telemetry for one successful network step inside a training session.
11///
12/// All references are valid only for the duration of
13/// [`TrainingObserver::on_step`]. The event does not grant mutable access to
14/// network state.
15#[derive(Debug, Clone, Copy)]
16pub struct TrainingStepEvent<'a> {
17    /// 0-based index of this example in the session batch.
18    pub step_index: usize,
19    /// Scalar reward from the [`crate::TrainingExample`].
20    pub reward: f32,
21    /// Neuromodulator state in effect after the step (borrowed from the network).
22    pub modulators: &'a NeuroModulators,
23    /// Indices of neurons that spiked on this step.
24    pub spike_indices: &'a [usize],
25    /// Number of examples processed so far, including this one.
26    pub steps_processed: usize,
27    /// Cumulative spike count across the session, including this step.
28    pub total_spikes: u64,
29}
30
31/// Receives one [`TrainingStepEvent`] after each successful network step.
32///
33/// Returning `Err` aborts the session before the next example is stepped. The
34/// failing step's network update has already been applied.
35///
36/// This trait is generic (statically dispatched). There is no `dyn` call on
37/// the session hot path.
38pub trait TrainingObserver {
39    /// Error type that aborts the session. Displayed in [`crate::TrainerError::Observer`].
40    type Error: core::fmt::Display;
41
42    /// Called once after a successful [`crate::PlasticityTrainer::train_step`].
43    ///
44    /// # Errors
45    ///
46    /// Any error aborts the session immediately; the next example is not processed.
47    fn on_step(&mut self, event: TrainingStepEvent<'_>) -> Result<(), Self::Error>;
48}
49
50/// Zero-sized observer that ignores every event.
51///
52/// `run_session` does not construct or dispatch events; this type exists so tests
53/// can exercise [`TrainingObserver`] without a runtime callback.
54#[cfg(test)]
55pub(crate) struct NoopObserver;
56
57#[cfg(test)]
58#[inline(never)]
59pub(crate) fn discard_step_event(
60    _event: TrainingStepEvent<'_>,
61) -> Result<(), core::convert::Infallible> {
62    Ok(())
63}
64
65#[cfg(test)]
66impl TrainingObserver for NoopObserver {
67    type Error = core::convert::Infallible;
68
69    fn on_step(&mut self, event: TrainingStepEvent<'_>) -> Result<(), Self::Error> {
70        discard_step_event(event)
71    }
72}
73
74impl<F, E> TrainingObserver for F
75where
76    F: FnMut(TrainingStepEvent<'_>) -> Result<(), E>,
77    E: core::fmt::Display,
78{
79    type Error = E;
80
81    fn on_step(&mut self, event: TrainingStepEvent<'_>) -> Result<(), Self::Error> {
82        self(event)
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use neuromod::NeuroModulators;
90
91    #[test]
92    fn event_is_copy_and_has_no_mutable_network_access() {
93        let mods = NeuroModulators::default();
94        let spikes = [0usize, 2];
95        let event = TrainingStepEvent {
96            step_index: 0,
97            reward: 0.25,
98            modulators: &mods,
99            spike_indices: &spikes,
100            steps_processed: 1,
101            total_spikes: 2,
102        };
103        let copy = event;
104        assert_eq!(copy.step_index, 0);
105        assert!((copy.reward - 0.25).abs() < 1e-6);
106        assert_eq!(copy.spike_indices, &spikes);
107        assert_eq!(copy.steps_processed, 1);
108        assert_eq!(copy.total_spikes, 2);
109        // `&NeuroModulators` / `&[usize]` only — mutating `copy.reward` cannot
110        // touch the network, and there is no `&mut` field to reach it.
111        let mut local = copy;
112        local.reward = 1.0;
113        assert!((local.reward - 1.0).abs() < 1e-6);
114        assert!((event.reward - 0.25).abs() < 1e-6);
115        assert!((mods.dopamine - 0.0).abs() < 1e-6);
116    }
117
118    #[test]
119    fn closure_observer_receives_event() {
120        let mods = NeuroModulators::default();
121        let spikes: &[usize] = &[];
122        let event = TrainingStepEvent {
123            step_index: 3,
124            reward: -0.1,
125            modulators: &mods,
126            spike_indices: spikes,
127            steps_processed: 4,
128            total_spikes: 0,
129        };
130        let mut seen = None;
131        let mut observer = |e: TrainingStepEvent<'_>| -> Result<(), &'static str> {
132            seen = Some(e.step_index);
133            Ok(())
134        };
135        observer.on_step(event).expect("closure observer");
136        assert_eq!(seen, Some(3));
137    }
138
139    #[test]
140    fn noop_observer_on_step_is_ok() {
141        let mods = NeuroModulators::default();
142        let spikes: &[usize] = &[];
143        let event = TrainingStepEvent {
144            step_index: 0,
145            reward: 0.0,
146            modulators: &mods,
147            spike_indices: spikes,
148            steps_processed: 1,
149            total_spikes: 0,
150        };
151        NoopObserver
152            .on_step(event)
153            .expect("noop observer cannot fail");
154    }
155}
plasticity_lab/observer.rs

plasticity_lab/
observer.rs
