Fix compute prefill_steps as policy_steps (#287)

* Fix compute prefill_steps as policy_steps

* Adjust prefill_steps once before training

sheeprl/algos/dreamer_v1/dreamer_v1.py

@@ -510,10 +510,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = (cfg.algo.learning_starts // policy_steps_per_iter) if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -644,7 +645,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 with timer("Time/train_time", SumMetric, sync_on_compute=cfg.metric.sync_on_compute):

sheeprl/algos/dreamer_v2/dreamer_v2.py

@@ -540,10 +540,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = cfg.algo.total_steps // policy_steps_per_iter if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -679,7 +680,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 local_data = rb.sample_tensors(

sheeprl/algos/dreamer_v3/dreamer_v3.py

@@ -508,10 +508,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // fabric.world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -657,7 +658,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 local_data = rb.sample_tensors(

sheeprl/algos/droq/droq.py

@@ -261,10 +261,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // fabric.world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -346,7 +347,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 train(

sheeprl/algos/p2e_dv1/p2e_dv1_exploration.py

@@ -534,10 +534,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = (cfg.algo.learning_starts // policy_steps_per_iter) if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -668,7 +669,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 with timer("Time/train_time", SumMetric, sync_on_compute=cfg.metric.sync_on_compute):

sheeprl/algos/p2e_dv1/p2e_dv1_finetuning.py

@@ -205,10 +205,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = (cfg.algo.learning_starts // policy_steps_per_iter) if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if resume_from_checkpoint:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -323,7 +324,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 if player.actor_type != "task":

sheeprl/algos/p2e_dv2/p2e_dv2_exploration.py

@@ -673,10 +673,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = cfg.algo.total_steps // policy_steps_per_iter if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -812,7 +813,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 local_data = rb.sample_tensors(

sheeprl/algos/p2e_dv2/p2e_dv2_finetuning.py

@@ -223,10 +223,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * world_size)
     total_iters = cfg.algo.total_steps // policy_steps_per_iter if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if resume_from_checkpoint:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -346,7 +347,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 if player.actor_type != "task":

sheeprl/algos/p2e_dv3/p2e_dv3_exploration.py

@@ -746,10 +746,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -895,7 +896,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 local_data = rb.sample_tensors(

sheeprl/algos/p2e_dv3/p2e_dv3_finetuning.py

@@ -210,10 +210,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if resume_from_checkpoint:
         cfg.algo.per_rank_batch_size = state["batch_size"] // world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -343,7 +344,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any], exploration_cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 if player.actor_type != "task":

sheeprl/algos/sac/sac.py

@@ -212,10 +212,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // fabric.world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)

sheeprl/algos/sac/sac_decoupled.py

@@ -145,9 +145,10 @@ def player(
     policy_steps_per_iter = int(cfg.env.num_envs)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from and not cfg.buffer.checkpoint:
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -227,7 +228,7 @@ def player(
 
         # Send data to the training agents
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / (fabric.world_size - 1))
             cumulative_per_rank_gradient_steps += per_rank_gradient_steps
             if per_rank_gradient_steps > 0:

sheeprl/algos/sac_ae/sac_ae.py

@@ -284,10 +284,11 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
     policy_steps_per_iter = int(cfg.env.num_envs * fabric.world_size)
     total_iters = int(cfg.algo.total_steps // policy_steps_per_iter) if not cfg.dry_run else 1
     learning_starts = cfg.algo.learning_starts // policy_steps_per_iter if not cfg.dry_run else 0
-    prefill_steps = learning_starts + start_iter
+    prefill_steps = learning_starts - int(learning_starts > 0)
     if cfg.checkpoint.resume_from:
         cfg.algo.per_rank_batch_size = state["batch_size"] // fabric.world_size
         learning_starts += start_iter
+        prefill_steps += start_iter
 
     # Create Ratio class
     ratio = Ratio(cfg.algo.replay_ratio, pretrain_steps=cfg.algo.per_rank_pretrain_steps)
@@ -374,7 +375,7 @@ def main(fabric: Fabric, cfg: Dict[str, Any]):
 
         # Train the agent
         if iter_num >= learning_starts:
-            ratio_steps = policy_step - prefill_steps + policy_steps_per_iter
+            ratio_steps = policy_step - prefill_steps * policy_steps_per_iter
             per_rank_gradient_steps = ratio(ratio_steps / world_size)
             if per_rank_gradient_steps > 0:
                 # We sample one time to reduce the communications between processes