[gym/rllib] Update checkpoint_dir so it is dependant on iteration num…

…ber (#789)
duburcqa · May 8, 2024 · 66113be · 66113be
1 parent 546b169
commit 66113be
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py b/python/gym_jiminy/rllib/gym_jiminy/rllib/utilities.py
@@ -493,7 +493,8 @@ def train(algo: Algorithm,
 
             # Backup the policy
             if checkpoint_period > 0 and iter_num % checkpoint_period == 0:
-                algo.save(checkpoint_dir=algo.logdir)
+                algo.save(
+                    checkpoint_dir=f"{algo.logdir}/checkpoint_{iter_num :06d}")
 
             # Check terminal conditions
             if 0 < max_timesteps < result["timesteps_total"]:
@@ -516,7 +517,8 @@ def train(algo: Algorithm,
         result_logger.close()
 
     # Backup trained agent and return file location
-    return algo.save(checkpoint_dir=algo.logdir)
+    return algo.save(
+        checkpoint_dir=f"{algo.logdir}/checkpoint_{iter_num :06d}")
 
 
 def _restore_default_connectors() -> None: