burn_train/learner/rl/
checkpointer.rs

1use burn_rl::{Policy, PolicyLearner, PolicyState};
2
3use crate::RLAgentRecord;
4use crate::{
5    RLComponentsTypes, RLPolicyRecord,
6    checkpoint::Checkpointer,
7    checkpoint::{AsyncCheckpointer, Checkpoint, CheckpointingAction, CheckpointingStrategy},
8    metric::store::EventStoreClient,
9};
10
11#[derive(new)]
12/// Used to create, delete, or load checkpoints of the training process.
13pub struct RLCheckpointer<RLC: RLComponentsTypes> {
14    policy: AsyncCheckpointer<RLPolicyRecord<RLC>>,
15    learning_agent: AsyncCheckpointer<RLAgentRecord<RLC>>,
16    strategy: Box<dyn CheckpointingStrategy>,
17}
18
19impl<RLC: RLComponentsTypes> RLCheckpointer<RLC>
20where
21    RLPolicyRecord<RLC>: Checkpoint,
22    RLAgentRecord<RLC>: Checkpoint,
23{
24    /// Create checkpoint for the training process.
25    pub fn checkpoint(
26        &mut self,
27        policy: &RLC::PolicyState,
28        learning_agent: &RLC::LearningAgent,
29        epoch: usize,
30        store: &EventStoreClient,
31    ) {
32        let actions = self.strategy.checkpointing(epoch, store);
33
34        for action in actions {
35            match action {
36                CheckpointingAction::Delete(epoch) => {
37                    self.policy
38                        .delete(epoch)
39                        .expect("Can delete policy checkpoint.");
40                    self.learning_agent
41                        .delete(epoch)
42                        .expect("Can delete learning agent checkpoint.")
43                }
44                CheckpointingAction::Save => {
45                    self.policy
46                        .save(epoch, policy.clone().into_record())
47                        .expect("Can save policy checkpoint.");
48                    self.learning_agent
49                        .save(epoch, learning_agent.record())
50                        .expect("Can save learning agent checkpoint.");
51                }
52            }
53        }
54    }
55
56    /// Load a training checkpoint.
57    pub fn load_checkpoint(
58        &self,
59        learning_agent: RLC::LearningAgent,
60        epoch: usize,
61    ) -> RLC::LearningAgent {
62        let record = self
63            .policy
64            .restore(epoch)
65            .expect("Can load model checkpoint.");
66        let policy = learning_agent.policy().load_record(record);
67
68        let record = self
69            .learning_agent
70            .restore(epoch)
71            .expect("Can load learning agent checkpoint.");
72        let mut learning_agent = learning_agent.load_record(record);
73        learning_agent.update_policy(policy);
74
75        learning_agent
76    }
77}
burn_train/learner/rl/checkpointer.rs

burn_train/learner/rl/
checkpointer.rs
