Update base_agent.py

xbpeng · Oct 19, 2024 · a7a9035 · a7a9035
1 parent c93d87c
commit a7a9035
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/learning/base_agent.py b/learning/base_agent.py
@@ -59,11 +59,12 @@ def train_model(self, max_samples, out_model_file, int_output_dir, log_file):
         while self._sample_count < max_samples:
             train_info = self._train_iter()
 
-            output_iter = (self._iter % self._iters_per_output == 0)
+            self._sample_count = self._update_sample_count()
+            output_iter = (self._iter % self._iters_per_output == 0) or (self._sample_count >= max_samples)
+
             if (output_iter):
                 test_info = self.test_model(self._test_episodes)
 
-            self._sample_count = self._update_sample_count()
             self._log_train_info(train_info, test_info, start_time) 
             self._logger.print_log()