burn_train/metric/processor/
rl_processor.rs

1use std::sync::Arc;
2
3use crate::{
4    EpisodeSummary, EvaluationItem, EventProcessorError, EventProcessorTraining, ItemLazy,
5    LearnerSummary, RLMetrics,
6    logger::TrainingProgressLogger,
7    metric::store::{Event, EventStoreClient, MetricsUpdate, Split},
8    renderer::{MetricState, MetricsRenderer},
9};
10
11/// Event happening during reinforcement learning.
12pub enum RLEvent<TS, ES> {
13    /// Signal the start of the process (e.g., learning starts).
14    Start {
15        /// The total number of items to process during training (e.g., total number of environment steps).
16        total_items: usize,
17        /// An optional label for this training.
18        label: Option<String>,
19    },
20    /// Signal an agent's training step.
21    TrainStep(EvaluationItem<TS>),
22    /// Signal a timestep of the agent-environment interface.
23    EnvStep(EvaluationItem<ES>),
24    /// Signal an episode end.
25    EpisodeEnd(EvaluationItem<EpisodeSummary>),
26    /// Signal the end of the process (e.g., learning ends).
27    End(Option<LearnerSummary>),
28}
29
30/// Event happening during evaluation of a reinforcement learning's agent.
31pub enum AgentEvaluationEvent<T> {
32    /// Signal the start of the process (e.g., training start)
33    Start(usize),
34    /// Signal a timestep of the agent-environment interface.
35    EnvStep(EvaluationItem<T>),
36    /// Signal an episode end.
37    EpisodeEnd(EvaluationItem<EpisodeSummary>),
38    /// Signal the end of the process (e.g., training end).
39    End,
40}
41
42/// An [event processor](EventProcessorTraining) that handles:
43///   - Computing and storing metrics in an [event store](crate::metric::store::EventStore).
44///   - Render metrics using a [metrics renderer](MetricsRenderer).
45pub struct RLEventProcessor<TS: ItemLazy, ES: ItemLazy> {
46    metrics: RLMetrics<TS, ES>,
47    renderer: Box<dyn MetricsRenderer>,
48    store: Arc<EventStoreClient>,
49    training_progress_logger: Option<Box<dyn TrainingProgressLogger>>,
50}
51
52impl<TS: ItemLazy, ES: ItemLazy> RLEventProcessor<TS, ES> {
53    pub(crate) fn new(
54        metrics: RLMetrics<TS, ES>,
55        renderer: Box<dyn MetricsRenderer>,
56        store: Arc<EventStoreClient>,
57    ) -> Self {
58        Self {
59            metrics,
60            renderer,
61            store,
62            training_progress_logger: None,
63        }
64    }
65
66    fn process_update_train(&mut self, update: MetricsUpdate) {
67        self.store
68            .add_event_train(crate::metric::store::Event::MetricsUpdate(update.clone()));
69
70        update
71            .entries
72            .into_iter()
73            .for_each(|entry| self.renderer.update_train(MetricState::Generic(entry)));
74
75        update
76            .entries_numeric
77            .into_iter()
78            .for_each(|numeric_update| {
79                self.renderer.update_train(MetricState::Numeric(
80                    numeric_update.entry,
81                    numeric_update.numeric_entry,
82                ))
83            });
84    }
85
86    fn process_update_valid(&mut self, update: MetricsUpdate) {
87        self.store
88            .add_event_valid(crate::metric::store::Event::MetricsUpdate(update.clone()));
89
90        update
91            .entries
92            .into_iter()
93            .for_each(|entry| self.renderer.update_valid(MetricState::Generic(entry)));
94
95        update
96            .entries_numeric
97            .into_iter()
98            .for_each(|numeric_update| {
99                self.renderer.update_valid(MetricState::Numeric(
100                    numeric_update.entry,
101                    numeric_update.numeric_entry,
102                ))
103            });
104    }
105}
106
107impl<TS: ItemLazy, ES: ItemLazy> EventProcessorTraining<RLEvent<TS, ES>, AgentEvaluationEvent<ES>>
108    for RLEventProcessor<TS, ES>
109{
110    fn process_train(&mut self, event: RLEvent<TS, ES>) -> Result<(), EventProcessorError> {
111        let mut failures = Vec::new();
112        match event {
113            RLEvent::Start { total_items, label } => {
114                let definitions = self.metrics.metric_definitions();
115                self.store
116                    .add_event_train(Event::MetricsInit(definitions.clone()));
117                definitions
118                    .iter()
119                    .for_each(|definition| self.renderer.register_metric(definition.clone()));
120                if let Some(logger) = &mut self.training_progress_logger {
121                    logger.start(0, 0, Some(total_items), label.as_deref());
122                }
123                self.renderer
124                    .start(0, 0, Some(total_items), label.as_deref());
125            }
126            RLEvent::TrainStep(item) => {
127                let item = match item.sync() {
128                    Ok(item) => item,
129                    Err(error) => return Err(EventProcessorError::sync(Split::Train, error)),
130                };
131                let metadata = (&item).into();
132
133                let (update, failed) = self.metrics.update_train_step(&item, &metadata);
134                failures.extend(failed);
135                self.process_update_train(update);
136
137                if let Some(logger) = &mut self.training_progress_logger {
138                    logger.log_event_training("TrainStep".to_string());
139                }
140                self.renderer.log_event_training("TrainStep".to_string());
141            }
142            RLEvent::EnvStep(item) => {
143                let item = match item.sync() {
144                    Ok(item) => item,
145                    Err(error) => return Err(EventProcessorError::sync(Split::Train, error)),
146                };
147                let metadata = (&item).into();
148
149                let (update, failed) = self.metrics.update_env_step(&item, &metadata);
150                failures.extend(failed);
151                self.process_update_train(update);
152
153                if let Some(logger) = &mut self.training_progress_logger {
154                    logger.update_split(item.progress.items_processed);
155                    logger.log_event_training("EnvStep".to_string());
156                }
157                self.renderer.update_split(item.progress.items_processed);
158                self.renderer.log_event_training("EnvStep".to_string());
159            }
160            RLEvent::EpisodeEnd(item) => {
161                let item = match item.sync() {
162                    Ok(item) => item,
163                    Err(error) => return Err(EventProcessorError::sync(Split::Train, error)),
164                };
165                let metadata = (&item).into();
166
167                let (update, failed) = self.metrics.update_episode_end(&item, &metadata);
168                failures.extend(failed);
169                self.process_update_train(update);
170
171                if let Some(logger) = &mut self.training_progress_logger {
172                    logger.log_event_training("EpisodeEnd".to_string());
173                }
174                self.renderer.log_event_training("EpisodeEnd".to_string());
175            }
176            RLEvent::End(learner_summary) => {
177                if let Some(logger) = &mut self.training_progress_logger {
178                    logger.end();
179                }
180                self.renderer.end();
181                self.renderer.on_train_end(learner_summary).ok();
182            }
183        }
184        EventProcessorError::from_errors(failures)
185    }
186
187    fn process_valid(
188        &mut self,
189        event: AgentEvaluationEvent<ES>,
190    ) -> Result<(), EventProcessorError> {
191        let mut failures = Vec::new();
192        match event {
193            AgentEvaluationEvent::Start(num_episodes) => {
194                if let Some(logger) = &mut self.training_progress_logger {
195                    logger.start_split("valid", num_episodes);
196                }
197                self.renderer.start_split("valid", num_episodes);
198            }
199            AgentEvaluationEvent::EnvStep(item) => {
200                let item = match item.sync() {
201                    Ok(item) => item,
202                    Err(error) => return Err(EventProcessorError::sync(Split::Valid, error)),
203                };
204                let metadata = (&item).into();
205
206                let (update, failed) = self.metrics.update_env_step_valid(&item, &metadata);
207                failures.extend(failed);
208                self.process_update_valid(update);
209
210                if let Some(logger) = &mut self.training_progress_logger {
211                    logger.log_event_training("EnvStep".to_string());
212                }
213                self.renderer.log_event_training("EnvStep".to_string());
214            }
215            AgentEvaluationEvent::EpisodeEnd(item) => {
216                let item = match item.sync() {
217                    Ok(item) => item,
218                    Err(error) => return Err(EventProcessorError::sync(Split::Valid, error)),
219                };
220                let metadata = (&item).into();
221
222                let (update, failed) = self.metrics.update_episode_end_valid(&item, &metadata);
223                failures.extend(failed);
224                self.process_update_valid(update);
225
226                if let Some(logger) = &mut self.training_progress_logger {
227                    logger.update_split(item.progress.items_processed);
228                    logger.log_event_training("EpisodeEnd".to_string());
229                }
230                self.renderer.update_split(item.progress.items_processed);
231                self.renderer.log_event_training("EpisodeEnd".to_string());
232            }
233            AgentEvaluationEvent::End => {
234                if let Some(logger) = &mut self.training_progress_logger {
235                    logger.end_split();
236                }
237                self.renderer.end_split();
238            }
239        }
240        EventProcessorError::from_errors(failures)
241    }
242
243    fn renderer(self) -> Box<dyn MetricsRenderer> {
244        self.renderer
245    }
246}
burn_train/metric/processor/rl_processor.rs

burn_train/metric/processor/
rl_processor.rs
