burn_train/learner/rl/
paradigm.rs

1use crate::checkpoint::{
2    AsyncCheckpointer, Checkpoint, CheckpointingStrategy, ComposedCheckpointingStrategy,
3    FileCheckpointer, KeepLastNCheckpoints, MetricCheckpointingStrategy,
4};
5use crate::learner::base::Interrupter;
6use crate::logger::{FileMetricLogger, MetricLogger};
7use crate::metric::store::{Aggregate, Direction, EventStoreClient, LogEventStore, Split};
8use crate::metric::{Adaptor, EpisodeLengthMetric, Metric, Numeric};
9use crate::renderer::{MetricsRenderer, default_renderer};
10use crate::{
11    ApplicationLoggerInstaller, AsyncProcessorTraining, FileApplicationLoggerInstaller, ItemLazy,
12    LearnerSummaryConfig, OffPolicyConfig, OffPolicyStrategy, RLAgentRecord, RLCheckpointer,
13    RLComponents, RLComponentsMarker, RLComponentsTypes, RLEventProcessor, RLMetrics,
14    RLPolicyRecord, RLStrategy,
15};
16use crate::{EpisodeSummary, RLStrategies};
17use burn_core::tensor::Device;
18use burn_rl::{
19    Batchable, Environment, EnvironmentInit, Policy, PolicyLearner, PolicyState, SliceAccess,
20    ToAction, ToObservation,
21};
22use std::collections::BTreeSet;
23use std::path::{Path, PathBuf};
24use std::sync::Arc;
25
26/// Structure to configure and launch reinforcement learning trainings.
27pub struct RLTraining<RLC: RLComponentsTypes> {
28    // Not that complex. Extracting into yet another type would only make it more confusing.
29    #[allow(clippy::type_complexity)]
30    checkpointers: Option<(
31        AsyncCheckpointer<RLPolicyRecord<RLC>>,
32        AsyncCheckpointer<RLAgentRecord<RLC>>,
33    )>,
34    num_steps: usize,
35    checkpoint: Option<usize>,
36    directory: PathBuf,
37    grad_accumulation: Option<usize>,
38    renderer: Option<Box<dyn MetricsRenderer + 'static>>,
39    metrics: RLMetrics<RLC::TrainingOutput, RLC::ActionContext>,
40    event_store: LogEventStore,
41    interrupter: Interrupter,
42    tracing_logger: Option<Box<dyn ApplicationLoggerInstaller>>,
43    checkpointer_strategy: Box<dyn CheckpointingStrategy>,
44    learning_strategy: RLStrategies<RLC>,
45    // Use BTreeSet instead of HashSet for consistent (alphabetical) iteration order
46    summary_metrics: BTreeSet<String>,
47    summary: bool,
48    env_initializer: RLC::EnvInit,
49    inference_device: Device,
50    label: Option<String>,
51}
52
53impl<E, EI, A> RLTraining<RLComponentsMarker<E, EI, A>>
54where
55    E: Environment + 'static,
56    EI: EnvironmentInit<E> + Send + 'static,
57    A: PolicyLearner + Send + 'static,
58    <A as PolicyLearner>::Record: Checkpoint,
59    A::TrainContext: ItemLazy + Clone + Send,
60    A::InnerPolicy: Policy + Send,
61    <A::InnerPolicy as Policy>::Observation: Batchable + Clone + Send,
62    <A::InnerPolicy as Policy>::ActionDistribution: Batchable + Clone + Send,
63    <A::InnerPolicy as Policy>::Action: Batchable + Clone + Send,
64    <A::InnerPolicy as Policy>::ActionContext: ItemLazy + Clone + Send + 'static,
65    <A::InnerPolicy as Policy>::PolicyState: Clone + Send,
66    <<A::InnerPolicy as Policy>::PolicyState as PolicyState>::Record: Checkpoint,
67    E::State: ToObservation<<A::InnerPolicy as Policy>::Observation> + Clone + Send + 'static,
68    E::Action: From<<A::InnerPolicy as Policy>::Action>
69        + ToAction<<A::InnerPolicy as Policy>::Action>
70        + Clone
71        + Send
72        + 'static,
73{
74    /// Creates a new runner for reinforcement learning.
75    ///
76    /// # Arguments
77    ///
78    /// * `directory` - The directory to save the checkpoints.
79    /// * `env_init` - Specifies how to initialize the environment.
80    pub fn new(directory: impl AsRef<Path>, env_initializer: EI) -> Self {
81        let directory = directory.as_ref().to_path_buf();
82        let experiment_log_file = directory.join("experiment.log");
83        Self {
84            num_steps: 1,
85            checkpoint: None,
86            checkpointers: None,
87            directory,
88            grad_accumulation: None,
89            metrics: RLMetrics::default(),
90            event_store: LogEventStore::default(),
91            renderer: None,
92            interrupter: Interrupter::new(),
93            tracing_logger: Some(Box::new(FileApplicationLoggerInstaller::new(
94                experiment_log_file,
95            ))),
96            checkpointer_strategy: Box::new(
97                ComposedCheckpointingStrategy::builder()
98                    .add(KeepLastNCheckpoints::new(2))
99                    .add(MetricCheckpointingStrategy::new(
100                        &EpisodeLengthMetric::new(), // default to evaluations' cumulative reward.
101                        Aggregate::Mean,
102                        Direction::Lowest,
103                        Split::Valid,
104                    ))
105                    .build(),
106            ),
107            learning_strategy: RLStrategies::OffPolicyStrategy(OffPolicyConfig::new()),
108            summary_metrics: BTreeSet::new(),
109            summary: false,
110            env_initializer,
111            inference_device: Default::default(),
112            label: None,
113        }
114    }
115}
116
117impl<RLC: RLComponentsTypes + 'static> RLTraining<RLC> {
118    /// Replace the default learning strategy (Off Policy learning) with the provided one.
119    ///
120    /// # Arguments
121    ///
122    /// * `training_strategy` - The training strategy.
123    pub fn with_learning_strategy(mut self, learning_strategy: RLStrategies<RLC>) -> Self {
124        self.learning_strategy = learning_strategy;
125        self
126    }
127
128    /// Replace the default metric loggers with the provided ones.
129    ///
130    /// # Arguments
131    ///
132    /// * `logger` - The training logger.
133    pub fn with_metric_logger<ML>(mut self, logger: ML) -> Self
134    where
135        ML: MetricLogger + 'static,
136    {
137        self.event_store.register_logger(logger);
138        self
139    }
140
141    /// Update the checkpointing_strategy.
142    pub fn with_checkpointing_strategy<CS: CheckpointingStrategy + 'static>(
143        mut self,
144        strategy: CS,
145    ) -> Self {
146        self.checkpointer_strategy = Box::new(strategy);
147        self
148    }
149
150    /// Replace the default CLI renderer with a custom one.
151    ///
152    /// # Arguments
153    ///
154    /// * `renderer` - The custom renderer.
155    pub fn renderer(mut self, renderer: Box<dyn MetricsRenderer + 'static>) -> Self {
156        self.renderer = Some(renderer);
157        self
158    }
159
160    /// Register numerical metrics for a training step of the agent.
161    pub fn metrics_train<Me: TrainMetricRegistration<RLC>>(self, metrics: Me) -> Self {
162        metrics.register(self)
163    }
164
165    /// Register textual metrics for a training step of the agent.
166    pub fn text_metrics_train<Me: TrainTextMetricRegistration<RLC>>(self, metrics: Me) -> Self {
167        metrics.register(self)
168    }
169
170    /// Register numerical metrics for each action of the agent.
171    pub fn metrics_agent<Me: AgentMetricRegistration<RLC>>(self, metrics: Me) -> Self {
172        metrics.register(self)
173    }
174
175    /// Register textual metrics for each action of the agent.
176    pub fn text_metrics_agent<Me: AgentTextMetricRegistration<RLC>>(self, metrics: Me) -> Self {
177        metrics.register(self)
178    }
179
180    /// Register numerical metrics for a completed episode.
181    pub fn metrics_episode<Me: EpisodeMetricRegistration<RLC>>(self, metrics: Me) -> Self {
182        metrics.register(self)
183    }
184
185    /// Register textual metrics for a completed episode.
186    pub fn text_metrics_episode<Me: EpisodeTextMetricRegistration<RLC>>(self, metrics: Me) -> Self {
187        metrics.register(self)
188    }
189
190    /// Register a textual metric for a training step.
191    pub fn text_metric_train<Me: Metric + 'static>(mut self, metric: Me) -> Self
192    where
193        RLC::TrainingOutput: Adaptor<Me::Input>,
194    {
195        self.metrics.register_text_metric_train(metric);
196        self
197    }
198
199    /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for a training step.
200    pub fn metric_train<Me>(mut self, metric: Me) -> Self
201    where
202        Me: Metric + Numeric + 'static,
203        RLC::TrainingOutput: Adaptor<Me::Input>,
204    {
205        self.summary_metrics.insert(metric.name().to_string());
206        self.metrics.register_metric_train(metric);
207        self
208    }
209
210    /// Register a textual metric for each action taken by the agent.
211    pub fn text_metric_agent<Me: Metric + 'static>(mut self, metric: Me) -> Self
212    where
213        RLC::ActionContext: Adaptor<Me::Input>,
214    {
215        self.metrics.register_text_metric_agent(metric.clone());
216        self.metrics.register_text_metric_agent_valid(metric);
217        self
218    }
219
220    /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for each action taken by the agent.
221    pub fn metric_agent<Me>(mut self, metric: Me) -> Self
222    where
223        Me: Metric + Numeric + 'static,
224        RLC::ActionContext: Adaptor<Me::Input>,
225    {
226        self.summary_metrics.insert(metric.name().to_string());
227        self.metrics.register_agent_metric(metric.clone());
228        self.metrics.register_agent_metric_valid(metric);
229        self
230    }
231
232    /// Register a textual metric for a completed episode.
233    pub fn text_metric_episode<Me: Metric + 'static>(mut self, metric: Me) -> Self
234    where
235        EpisodeSummary: Adaptor<Me::Input> + 'static,
236    {
237        self.metrics.register_text_metric_episode(metric.clone());
238        self.metrics.register_text_metric_episode_valid(metric);
239        self
240    }
241
242    /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for a completed episode.
243    pub fn metric_episode<Me>(mut self, metric: Me) -> Self
244    where
245        Me: Metric + Numeric + 'static,
246        EpisodeSummary: Adaptor<Me::Input> + 'static,
247    {
248        self.summary_metrics.insert(metric.name().to_string());
249        self.metrics.register_episode_metric(metric.clone());
250        self.metrics.register_episode_metric_valid(metric);
251        self
252    }
253
254    /// The number of environment steps to train for.
255    pub fn num_steps(mut self, num_steps: usize) -> Self {
256        self.num_steps = num_steps;
257        self
258    }
259
260    /// The step from which the training must resume.
261    pub fn checkpoint(mut self, checkpoint: usize) -> Self {
262        self.checkpoint = Some(checkpoint);
263        self
264    }
265
266    /// Provides a handle that can be used to interrupt training.
267    pub fn interrupter(&self) -> Interrupter {
268        self.interrupter.clone()
269    }
270
271    /// Override the handle for stopping training with an externally provided handle
272    pub fn with_interrupter(mut self, interrupter: Interrupter) -> Self {
273        self.interrupter = interrupter;
274        self
275    }
276
277    /// By default, Rust logs are captured and written into
278    /// `experiment.log`. If disabled, standard Rust log handling
279    /// will apply.
280    pub fn with_application_logger(
281        mut self,
282        logger: Option<Box<dyn ApplicationLoggerInstaller>>,
283    ) -> Self {
284        self.tracing_logger = logger;
285        self
286    }
287
288    /// Register a checkpointer that will save the environment runner's [policy](Policy)
289    /// and the [PolicyLearner](PolicyLearner) state to different files.
290    pub fn with_checkpointer(mut self) -> Self
291    where
292        RLPolicyRecord<RLC>: Checkpoint,
293        RLAgentRecord<RLC>: Checkpoint,
294    {
295        let checkpoint_dir = self.directory.join("checkpoint");
296        let checkpointer_policy = FileCheckpointer::new(&checkpoint_dir, "policy");
297        let checkpointer_learning = FileCheckpointer::new(&checkpoint_dir, "learning-agent");
298
299        self.checkpointers = Some((
300            AsyncCheckpointer::new(checkpointer_policy),
301            AsyncCheckpointer::new(checkpointer_learning),
302        ));
303
304        self
305    }
306
307    /// The device on which to run inference during rollout collection and validation.
308    pub fn with_inference_device(mut self, device: Device) -> Self {
309        self.inference_device = device;
310        self
311    }
312
313    /// Enable the training summary report.
314    ///
315    /// The summary will be displayed after `.launch()`, when the renderer is dropped.
316    pub fn summary(mut self) -> Self {
317        self.summary = true;
318        self
319    }
320
321    /// Set a label for this training, making it easier to differentiate multiple runs.
322    pub fn label(mut self, label: &str) -> Self {
323        self.label = Some(label.to_string());
324        self
325    }
326
327    /// Launch the training with the specified [PolicyLearner](PolicyLearner) on the specified environment.
328    pub fn launch(mut self, learner_agent: RLC::LearningAgent) -> RLResult<RLC::Policy>
329    where
330        RLC::PolicyObs: SliceAccess,
331        RLC::PolicyAction: SliceAccess,
332    {
333        if self.tracing_logger.is_some()
334            && let Err(e) = self.tracing_logger.as_ref().unwrap().install()
335        {
336            log::warn!("Failed to install the experiment logger: {e}");
337        }
338        let renderer = self
339            .renderer
340            .unwrap_or_else(|| default_renderer(self.interrupter.clone(), self.checkpoint));
341
342        if !self.event_store.has_loggers() {
343            self.event_store
344                .register_logger(FileMetricLogger::new(self.directory.clone()));
345        }
346
347        let event_store = Arc::new(EventStoreClient::new(self.event_store));
348        let event_processor = AsyncProcessorTraining::new(RLEventProcessor::new(
349            self.metrics,
350            renderer,
351            event_store.clone(),
352        ));
353
354        let checkpointer = self.checkpointers.map(|(policy, learning_agent)| {
355            RLCheckpointer::new(policy, learning_agent, self.checkpointer_strategy)
356        });
357
358        let summary = if self.summary {
359            Some(LearnerSummaryConfig {
360                directory: self.directory,
361                metrics: self.summary_metrics.into_iter().collect::<Vec<_>>(),
362            })
363        } else {
364            None
365        };
366
367        let components = RLComponents::<RLC> {
368            checkpoint: self.checkpoint,
369            checkpointer,
370            interrupter: self.interrupter,
371            event_processor,
372            event_store,
373            num_steps: self.num_steps,
374            grad_accumulation: self.grad_accumulation,
375            summary,
376            inference_device: self.inference_device,
377            label: self.label,
378        };
379
380        let mut learner_agent = learner_agent;
381        if let Some(checkpoint) = components.checkpoint
382            && let Some(checkpointer) = &components.checkpointer
383        {
384            learner_agent = checkpointer.load_checkpoint(learner_agent, checkpoint);
385        }
386
387        match self.learning_strategy {
388            RLStrategies::OffPolicyStrategy(config) => {
389                let strategy = OffPolicyStrategy::new(config);
390                strategy.train(learner_agent, components, self.env_initializer)
391            }
392            RLStrategies::Custom(strategy) => {
393                strategy.train(learner_agent, components, self.env_initializer)
394            }
395        }
396    }
397}
398
399/// The result of reinforcement learning, containing the final policy along with the [renderer](MetricsRenderer).
400pub struct RLResult<P> {
401    /// The learned policy.
402    pub policy: P,
403    /// The renderer that can be used for follow up training and evaluation.
404    pub renderer: Box<dyn MetricsRenderer>,
405    /// The stop that ended training early, if
406    /// [`Interrupter::stop`](crate::Interrupter::stop) was called and no error happened.
407    pub interrupted: Option<crate::Interruption>,
408    /// The error that ended training early, if it hit one.
409    pub error: Option<std::sync::Arc<crate::TrainingError>>,
410}
411
412/// Trait to fake variadic generics for train step metrics.
413pub trait AgentMetricRegistration<RLC: RLComponentsTypes>: Sized {
414    /// Register the metrics.
415    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
416}
417
418/// Trait to fake variadic generics for train step text metrics.
419pub trait AgentTextMetricRegistration<RLC: RLComponentsTypes>: Sized {
420    /// Register the metrics.
421    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
422}
423
424/// Trait to fake variadic generics for env step metrics.
425pub trait TrainMetricRegistration<RLC: RLComponentsTypes>: Sized {
426    /// Register the metrics.
427    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
428}
429
430/// Trait to fake variadic generics for env step text metrics.
431pub trait TrainTextMetricRegistration<RLC: RLComponentsTypes>: Sized {
432    /// Register the metrics.
433    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
434}
435
436/// Trait to fake variadic generics for episode metrics.
437pub trait EpisodeMetricRegistration<RLC: RLComponentsTypes>: Sized {
438    /// Register the metrics.
439    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
440}
441
442/// Trait to fake variadic generics for episode text metrics.
443pub trait EpisodeTextMetricRegistration<RLC: RLComponentsTypes>: Sized {
444    /// Register the metrics.
445    fn register(self, builder: RLTraining<RLC>) -> RLTraining<RLC>;
446}
447
448macro_rules! gen_tuple {
449    ($($M:ident),*) => {
450        impl<$($M,)* RLC: RLComponentsTypes + 'static> TrainTextMetricRegistration<RLC> for ($($M,)*)
451        where
452            $(RLC::TrainingOutput: Adaptor<$M::Input>,)*
453            $($M: Metric + 'static,)*
454        {
455            #[allow(non_snake_case)]
456            fn register(
457                self,
458                builder: RLTraining<RLC>,
459            ) -> RLTraining<RLC> {
460                let ($($M,)*) = self;
461                $(let builder = builder.text_metric_train($M.clone());)*
462                builder
463            }
464        }
465
466        impl<$($M,)* RLC: RLComponentsTypes + 'static> TrainMetricRegistration<RLC> for ($($M,)*)
467        where
468            $(RLC::TrainingOutput: Adaptor<$M::Input>,)*
469            $($M: Metric + Numeric + 'static,)*
470        {
471            #[allow(non_snake_case)]
472            fn register(
473                self,
474                builder: RLTraining<RLC>,
475            ) -> RLTraining<RLC> {
476                let ($($M,)*) = self;
477                $(let builder = builder.metric_train($M.clone());)*
478                builder
479            }
480        }
481
482        impl<$($M,)* RLC: RLComponentsTypes + 'static> AgentTextMetricRegistration<RLC> for ($($M,)*)
483        where
484            $(RLC::ActionContext: Adaptor<$M::Input>,)*
485            $($M: Metric + 'static,)*
486        {
487            #[allow(non_snake_case)]
488            fn register(
489                self,
490                builder: RLTraining<RLC>,
491            ) -> RLTraining<RLC> {
492                let ($($M,)*) = self;
493                $(let builder = builder.text_metric_agent($M.clone());)*
494                builder
495            }
496        }
497
498        impl<$($M,)* RLC: RLComponentsTypes + 'static> AgentMetricRegistration<RLC> for ($($M,)*)
499        where
500            $(RLC::ActionContext: Adaptor<$M::Input>,)*
501            $($M: Metric + Numeric + 'static,)*
502        {
503            #[allow(non_snake_case)]
504            fn register(
505                self,
506                builder: RLTraining<RLC>,
507            ) -> RLTraining<RLC> {
508                let ($($M,)*) = self;
509                $(let builder = builder.metric_agent($M.clone());)*
510                builder
511            }
512        }
513
514        impl<$($M,)* RLC: RLComponentsTypes + 'static> EpisodeTextMetricRegistration<RLC> for ($($M,)*)
515        where
516            $(EpisodeSummary: Adaptor<$M::Input> + 'static,)*
517            $($M: Metric + 'static,)*
518        {
519            #[allow(non_snake_case)]
520            fn register(
521                self,
522                builder: RLTraining<RLC>,
523            ) -> RLTraining<RLC> {
524                let ($($M,)*) = self;
525                $(let builder = builder.text_metric_episode($M.clone());)*
526                builder
527            }
528        }
529
530        impl<$($M,)* RLC: RLComponentsTypes + 'static> EpisodeMetricRegistration<RLC> for ($($M,)*)
531        where
532            $(EpisodeSummary: Adaptor<$M::Input> + 'static,)*
533            $($M: Metric + Numeric + 'static,)*
534        {
535            #[allow(non_snake_case)]
536            fn register(
537                self,
538                builder: RLTraining<RLC>,
539            ) -> RLTraining<RLC> {
540                let ($($M,)*) = self;
541                $(let builder = builder.metric_episode($M.clone());)*
542                builder
543            }
544        }
545    };
546}
547
548gen_tuple!(M1);
549gen_tuple!(M1, M2);
550gen_tuple!(M1, M2, M3);
551gen_tuple!(M1, M2, M3, M4);
552gen_tuple!(M1, M2, M3, M4, M5);
553gen_tuple!(M1, M2, M3, M4, M5, M6);
burn_train/learner/rl/paradigm.rs

burn_train/learner/rl/
paradigm.rs
