JAX PPO implementation incompatible with reverb 0.7.1

[kinalmehta]

JAX ecosystem seems to be incompatible with reverb 0.7.1 and raises the below error

Traceback (most recent call last):
  File "/home/kinal/Desktop/marl/acme/examples/gym/run_ppo.py", line 99, in <module>
    app.run(main)
  File "/home/kinal/miniconda3/envs/dm-env/lib/python3.9/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/kinal/miniconda3/envs/dm-env/lib/python3.9/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/home/kinal/Desktop/marl/acme/examples/gym/run_ppo.py", line 61, in main
    agent = ppo.PPO(
  File "/home/kinal/Desktop/marl/acme/acme/agents/jax/ppo/agents.py", line 120, in __init__
    super().__init__(
  File "/home/kinal/Desktop/marl/acme/acme/jax/layouts/local_layout.py", line 105, in __init__
    rate_limiter = reverb.rate_limiters.RateLimiter(
TypeError: __init__() got an unexpected keyword argument 'samples_per_insert'

Going through the reverb repo it seems to be compatible with the latest commits. Any suggestions on how to solve this?

Reproduction script

python examples/gym/run_ppo.py

[KaleabTessera]

@kinalmehta There is a bug in reverb 0.7.1, where you can't set the rate limiters. This was fixed a few days back (https://github.com/deepmind/reverb/commit/9d6e120c2638a2d8ef450410eae2f6a664c0a69d).

Using the nightly version of reverb and launchpad worked for me:

pip install dm-reverb-nightly[tensorflow]
pip install dm-launchpad-nightly[tensorflow]