burn_train/learner/rl/
strategy.rs

1use std::sync::Arc;
2
3use burn_core::tensor::Device;
4
5use crate::{
6    Interrupter, LearnerSummaryConfig, OffPolicyConfig, RLCheckpointer, RLComponentsTypes, RLEvent,
7    RLEventProcessorType, RLResult,
8    metric::{processor::EventProcessorTraining, store::EventStoreClient},
9};
10
11/// Struct to minimise parameters passed to [RLStrategy::train].
12pub struct RLComponents<RLC: RLComponentsTypes> {
13    /// An optional label for this training.
14    pub label: Option<String>,
15    /// The total number of environment steps.
16    pub num_steps: usize,
17    /// The step number from which to continue the training.
18    pub checkpoint: Option<usize>,
19    /// A checkpointer used to load and save learning checkpoints.
20    pub checkpointer: Option<RLCheckpointer<RLC>>,
21    /// Enables gradients accumulation.
22    pub grad_accumulation: Option<usize>,
23    /// An [Interupter](Interrupter) that allows aborting the training/evaluation process early.
24    pub interrupter: Interrupter,
25    /// An [EventProcessor](crate::EventProcessorTraining) that processes events happening during training and evaluation.
26    pub event_processor: RLEventProcessorType<RLC>,
27    /// A reference to an [EventStoreClient](EventStoreClient).
28    pub event_store: Arc<EventStoreClient>,
29    /// Config for creating a summary of the learning
30    pub summary: Option<LearnerSummaryConfig>,
31    /// Device used for running inference during environmment sampling or validation.
32    pub inference_device: Device,
33}
34
35/// The strategy for reinforcement learning.
36#[derive(Clone)]
37pub enum RLStrategies<RLC: RLComponentsTypes> {
38    /// Training on one device
39    OffPolicyStrategy(OffPolicyConfig),
40    /// Training using a custom learning strategy
41    Custom(CustomRLStrategy<RLC>),
42}
43
44/// A reference to an implementation of [RLStrategy].
45pub type CustomRLStrategy<LC> = Arc<dyn RLStrategy<LC>>;
46
47/// Provides the `fit` function for any learning strategy
48pub trait RLStrategy<RLC: RLComponentsTypes> {
49    /// Train the learner agent with this strategy.
50    fn train(
51        &self,
52        mut learner_agent: RLC::LearningAgent,
53        mut training_components: RLComponents<RLC>,
54        env_init: RLC::EnvInit,
55    ) -> RLResult<RLC::Policy> {
56        let starting_epoch = training_components.checkpoint.unwrap_or(0) + 1;
57        let summary_config = training_components.summary.clone();
58        let interrupter = training_components.interrupter.clone();
59
60        // Event processor start training
61        interrupter.fail_on_error(training_components.event_processor.process_train(
62            RLEvent::Start {
63                total_items: training_components.num_steps,
64                label: training_components.label.clone(),
65            },
66        ));
67
68        // Training loop
69        let (policy, mut event_processor) = self.train_loop(
70            training_components,
71            &mut learner_agent,
72            starting_epoch,
73            env_init,
74        );
75
76        let summary = summary_config.and_then(|summary| summary.init().ok());
77
78        // Signal training end. For the TUI renderer, this handles the exit & return to main screen.
79        // TODO: summary makes sense for RL?
80        interrupter.fail_on_error(event_processor.process_train(RLEvent::End(summary)));
81
82        // let model = model.valid();
83
84        // Finish processing remaining events.
85        interrupter.fail_on_error(event_processor.flush());
86        let renderer = event_processor.renderer();
87
88        RLResult {
89            policy,
90            renderer,
91            interrupted: interrupter.interruption(),
92            error: interrupter.error(),
93        }
94    }
95
96    /// Training loop for this strategy
97    fn train_loop(
98        &self,
99        training_components: RLComponents<RLC>,
100        learner_agent: &mut RLC::LearningAgent,
101        starting_epoch: usize,
102        env_init: RLC::EnvInit,
103    ) -> (RLC::Policy, RLEventProcessorType<RLC>);
104}
burn_train/learner/rl/strategy.rs

burn_train/learner/rl/
strategy.rs
