From 924adcf4024ec2fec85bc30e01aa0ca77d087b79 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Mon, 2 May 2022 18:30:46 +0200 Subject: [PATCH] [RLlib] Issue 24074: multi-GPU learner thread key error in MA-scenarios. (#24382) --- rllib/execution/multi_gpu_learner_thread.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/rllib/execution/multi_gpu_learner_thread.py b/rllib/execution/multi_gpu_learner_thread.py index a9c2848e6..61e33388a 100644 --- a/rllib/execution/multi_gpu_learner_thread.py +++ b/rllib/execution/multi_gpu_learner_thread.py @@ -209,12 +209,16 @@ class _MultiGPULoaderThread(threading.Thread): if not s.local_worker.is_policy_to_train(pid, batch): continue policy = policy_map[pid] - policy.load_batch_into_buffer( - batch=batch - if isinstance(batch, SampleBatch) - else batch.policy_batches[pid], - buffer_index=buffer_idx, - ) + if isinstance(batch, SampleBatch): + policy.load_batch_into_buffer( + batch=batch, + buffer_index=buffer_idx, + ) + elif pid in batch.policy_batches: + policy.load_batch_into_buffer( + batch=batch.policy_batches[pid], + buffer_index=buffer_idx, + ) # Tag just-loaded stack as "ready". s.ready_tower_stacks.put(buffer_idx)