[RLlib] Fix IMPALA/APPO when using multi GPU setup and Multi-Agent Env (

ray-project#35120) Signed-off-by: Michael <[email protected]> Co-authored-by: Artur Niederfahrenhorst <[email protected]>
scv119 · Jun 11, 2023 · 85d0140 · 85d0140
1 parent 7151cc7
commit 85d0140
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/rllib/BUILD b/rllib/BUILD
@@ -2713,6 +2713,16 @@ py_test(
  args = ["TestTrainAndEvaluate"]
 )
 
+py_test(
+ name = "tests/test_supported_multi_agent_multi_gpu",
+ main = "tests/test_supported_multi_agent.py",
+ tags = ["team:rllib", "tests_dir", "multi_gpu"],
+ size = "medium",
+ srcs = ["tests/test_supported_multi_agent.py"],
+ args = ["TestSupportedMultiAgentMultiGPU"]
+)
+
+
 py_test(
  name = "tests/test_supported_multi_agent_pg",
  main = "tests/test_supported_multi_agent.py",

diff --git a/rllib/execution/multi_gpu_learner_thread.py b/rllib/execution/multi_gpu_learner_thread.py
@@ -162,7 +162,9 @@ def step(self) -> None:
  default_policy_results = policy.learn_on_loaded_batch(
  offset=0, buffer_index=buffer_idx
  )
- learner_info_builder.add_learn_on_batch_results(default_policy_results)
+ learner_info_builder.add_learn_on_batch_results(
+ default_policy_results, policy_id=pid
+ )
  self.policy_ids_updated.append(pid)
  get_num_samples_loaded_into_buffer += (
  policy.get_num_samples_loaded_into_buffer(buffer_idx)

diff --git a/rllib/tests/test_supported_multi_agent.py b/rllib/tests/test_supported_multi_agent.py
@@ -160,6 +160,19 @@ def test_sac_multiagent(self):
  )
 
 
+class TestSupportedMultiAgentMultiGPU(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls) -> None:
+ ray.init()
+
+ @classmethod
+ def tearDownClass(cls) -> None:
+ ray.shutdown()
+
+ def test_impala_multiagent_multi_gpu(self):
+ check_support_multiagent("IMPALA", {"num_gpus": 2})
+
+
 if __name__ == "__main__":
  import pytest
  import sys