From 2d86091be1ccc73d3b81f2ad986139b97bbdc713 Mon Sep 17 00:00:00 2001 From: Avnish Narayan <38871737+avnishn@users.noreply.github.com> Date: Fri, 17 Mar 2023 08:55:36 -0700 Subject: [PATCH] [RLlib] Add option for running multiple sgd iters for impala learner api (#33316) Signed-off-by: Avnish Signed-off-by: Jack He --- rllib/algorithms/impala/impala.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/rllib/algorithms/impala/impala.py b/rllib/algorithms/impala/impala.py index 006a6f726807..cfd5f5bc3a74 100644 --- 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