[RLlib] Add option for running multiple sgd iters for impala learner …

…api (ray-project#33316) Signed-off-by: Avnish <[email protected]> Signed-off-by: Jack He <[email protected]>
ProjectsByJackHe · Mar 21, 2023 · 2d86091 · 2d86091
1 parent 9834dfe
commit 2d86091
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/rllib/algorithms/impala/impala.py b/rllib/algorithms/impala/impala.py
@@ -862,7 +862,10 @@ def learn_on_processed_samples(self) -> ResultDict:
             # Then we can't do async updates, so we need to block.
             blocking = self.config.num_learner_workers == 0
             lg_results = self.learner_group.update(
-                batch, reduce_fn=_reduce_impala_results, block=blocking
+                batch,
+                reduce_fn=_reduce_impala_results,
+                block=blocking,
+                num_iters=self.config.num_sgd_iter,
             )
         else:
             lg_results = None