Don't use the keys returned by JAX.
These keys have been cast to 32-bit precision and lose information. When sent to the Reverb server the server is unable to find the original key because of this. The solution is to use the original uint64 keys.
@@ -167,6 +167,10 @@ class DQNLearner(acme.Learner, acme.Saveable):
# Update priorities in replay.
if self._replay_client:
+ # Use samples.info.key instead of outputs.keys because JAX casts uint64
+ # to int32, which loses information.
+ outputs = LearnerOutputs(
+ keys=samples.info.key, priorities=outputs.priorities)
self._async_priority_updater.put(outputs)
# Write to logs.
