refact : save model

bot/src/burnrl/main.rs

@@ -1,13 +1,35 @@
 use bot::burnrl::{dqn_model, environment, utils::demo_model};
 use burn::backend::{Autodiff, NdArray};
+use burn::module::Module;
+use burn::record::{CompactRecorder, Recorder};
+use burn_rl::agent::DQN;
 use burn_rl::base::ElemType;
 
 type Backend = Autodiff<NdArray<ElemType>>;
 type Env = environment::TrictracEnvironment;
 
 fn main() {
+    println!("> Entraînement");
     let num_episodes = 3;
     let agent = dqn_model::run::<Env, Backend>(num_episodes, false); //true);
+    println!("> Sauvegarde");
+    save(&agent);
 
-    demo_model::<Env>(agent);
+    // cette ligne sert à extraire le "cerveau" de l'agent entraîné,
+    // sans les données nécessaires à l'entraînement
+    let valid_agent = agent.valid();
+
+    println!("> Test");
+    demo_model::<Env>(valid_agent);
+}
+
+fn save(agent: &DQN<Env, Backend, dqn_model::Net<Backend>>) {
+    let path = "models/burn_dqn".to_string();
+    let inference_network = agent.model().clone().into_record();
+    let recorder = CompactRecorder::new();
+    let model_path = format!("{}_model.burn", path);
+    println!("Modèle sauvegardé : {}", model_path);
+    recorder
+        .record(inference_network, model_path.into())
+        .unwrap();
 }