Call mutate_priorities() with the entire batch instead of looping.
Previously an RPC request was being sent for each key in the batch, which
slowed the learners down significantly. This change will send the entire batch of updates together as a single RPC
request.
@@ -124,9 +124,9 @@ class DQNLearner(acme.Learner, acme.Saveable):
return new_state, outputs
def update_priorities(outputs: LearnerOutputs):
- for key, priority in zip(outputs.keys, outputs.priorities):
replay_client.mutate_priorities(
- table=adders.DEFAULT_PRIORITY_TABLE, updates={key: priority})
+ table=adders.DEFAULT_PRIORITY_TABLE,
+ updates=dict(zip(outputs.keys, outputs.priorities)))
# Internalise agent components (replay buffer, networks, optimizer).
self._replay_client = replay_client
