Merge remote-tracking branch 'origin/dev' into feature/1572-fix-docs-formatting

src/primaite/agents/sb3.py

@@ -71,16 +71,19 @@ class SB3Agent(AgentSessionABC):
             verbose=self.sb3_output_verbose_level,
             n_steps=self._training_config.num_steps,
             tensorboard_log=str(self._tensorboard_log_path),
+            seed=self._training_config.seed,
         )
 
     def _save_checkpoint(self):
         checkpoint_n = self._training_config.checkpoint_every_n_episodes
         episode_count = self._env.episode_count
-        if checkpoint_n > 0 and episode_count > 0:
-            if (episode_count % checkpoint_n == 0) or (episode_count == self._training_config.num_episodes):
-                checkpoint_path = self.checkpoints_path / f"sb3ppo_{episode_count}.zip"
-                self._agent.save(checkpoint_path)
-                _LOGGER.debug(f"Saved agent checkpoint: {checkpoint_path}")
+        save_checkpoint = False
+        if checkpoint_n:
+            save_checkpoint = episode_count % checkpoint_n == 0
+        if episode_count and save_checkpoint:
+            checkpoint_path = self.checkpoints_path / f"sb3ppo_{episode_count}.zip"
+            self._agent.save(checkpoint_path)
+            _LOGGER.debug(f"Saved agent checkpoint: {checkpoint_path}")
 
     def _get_latest_checkpoint(self):
         pass
@@ -102,25 +105,27 @@ class SB3Agent(AgentSessionABC):
             self._agent.learn(total_timesteps=time_steps)
             self._save_checkpoint()
         self._env.reset()
+        self.save()
         self._env.close()
         super().learn()
 
+        # save agent
+        self.save()
+
     def evaluate(
         self,
-        deterministic: bool = True,
         **kwargs,
     ):
         """
         Evaluate the agent.
 
-        :param deterministic: Whether the evaluation is deterministic.
         :param kwargs: Any agent-specific key-word args to be passed.
         """
         time_steps = self._training_config.num_steps
         episodes = self._training_config.num_episodes
         self._env.set_as_eval()
         self.is_eval = True
-        if deterministic:
+        if self._training_config.deterministic:
             deterministic_str = "deterministic"
         else:
             deterministic_str = "non-deterministic"
@@ -131,7 +136,7 @@ class SB3Agent(AgentSessionABC):
             obs = self._env.reset()
 
             for step in range(time_steps):
-                action, _states = self._agent.predict(obs, deterministic=deterministic)
+                action, _states = self._agent.predict(obs, deterministic=self._training_config.deterministic)
                 if isinstance(action, np.ndarray):
                     action = np.int64(action)
                 obs, rewards, done, info = self._env.step(action)
@@ -146,7 +151,7 @@ class SB3Agent(AgentSessionABC):
 
     def save(self):
         """Save the agent."""
-        raise NotImplementedError
+        self._agent.save(self._saved_agent_path)
 
     def export(self):
         """Export the agent to transportable file format."""