Converting Tensorflow Dataset to `iterator` does not sync well with client

Hi, I accidentally stumbled upon a problem within the [tutorial notebook](https://colab.research.google.com/github/deepmind/acme/blob/master/examples/tutorial.ipynb) when playing around with the acme and reverb API that causes a weird synchronization behaviour between sampling from the `reverb` table and updating priorities. Another artifact of this that I encountered is that the very first transition would be consistently repeated until some hidden tensorflow buffer would be flushed.

What I found is that when I would mutate the priorities in a `reverb` table using `client.mutate_priorities(table_name, my_dict)` and then create an iterator from the `tf.data.Dataset` object, then the priorities would update only after flushing a large number of samples. In contrast, if I didn't convert the `tf.data.Dataset` to an iterator and used the `dataset.batch(n); dataset.take(n)` interface, it would immediately sync with the new priorities.

It seems to me that the problem lies with the implementation of `__iter__` in tf.data.Dataset, but I posted this issue here since the Colab makes a call to `as_numpy_iterator()` on the dataset object, and this is also the implementation of the `D4PG` jax agent. Since this is a silent and obscure bug, this effectively eliminates the possibility of changing the baseline `D4PG` agent to utilize Prioritized Experience Replay.

## Minimal Reproducible example:
```python
import warnings
warnings.filterwarnings('ignore')

import acme

from acme import wrappers
from acme.datasets import reverb as datasets
from acme.adders.reverb import sequence
from acme.jax import utils

import tree
import reverb
import jax

import numpy as np

from dm_control import suite

# Create dummy environment with short episodes to easily dichotomize samples
env = suite.load('cartpole', 'balance')
env = wrappers.step_limit.StepLimitWrapper(env, step_limit=5)
spec = acme.make_environment_spec(env)

# Danger: reverb.Table crashes kernel if run > once
table = reverb.Table(
name='priority_table',
sampler=reverb.selectors.Prioritized(priority_exponent=0.8),
remover=reverb.selectors.Fifo(),
max_size=10_000,
rate_limiter=reverb.rate_limiters.MinSize(1),
signature=sequence.SequenceAdder.signature(spec)
)

server = reverb.Server([table], port=None)
client = reverb.Client(f'localhost:{server.port}')

# Construct adder such that only 1 sample is added to table after an episode.
adder = sequence.SequenceAdder(client, sequence_length=6, period=5)

def new_dataset():
# Clear old data
client.reset(table.name)
return datasets.make_reverb_dataset(
table=table.name, server_address=client.server_address, batch_size=3
)

def fill_dataset():
step = env.reset()
adder.add_first(step)

action = env.action_spec().generate_value()
i = 0
while (not step.last()) and i < 10:
step = env.step(action)
adder.add(action, step)
i += 1

env.close()
adder.reset()


### Example of expected behaviour
dataset = new_dataset()
fill_dataset()

print('before mutation')
for s in dataset.take(1):
k, p = s.info.key.numpy().ravel(), s.info.priority.numpy().ravel()

print(s.data.action.numpy().reshape(3, -1)) # (B, T, 1) -> (B, T)
print('sample priority:', p)

# Iteratively halve the priorities
new_priorities = dict(zip(k, p * 0.5))
client.mutate_priorities(table.name, new_priorities)

print()

print('after mutation')
for s in dataset.take(1):
# Priorities have been updated --> all probabilities should now be adjusted.

print(s.data.action.numpy().reshape(3, -1)) # (B, T, 1) -> (B, T)
print('sample priority:', s.info.priority.numpy())

### Test-cases

print('\nUsing dataset.take')
dataset = new_dataset()
fill_dataset()

# This runs fine
for repeat in range(5):
for i in range(30): # Flush count guess
for s in dataset.take(1):
k, p = s.info.key.numpy().ravel(), s.info.priority.numpy().ravel()

# Exponentially decay the priorities
new_priorities = dict(zip(k, p * 0.999))
client.mutate_priorities(table.name, new_priorities)

for s in dataset.take(1):
new_p = s.info.priority.numpy().ravel()
assert not np.isclose(new_p, p).any(), "priorities did not update!"
else:
# No break in for loop
print('No errors!')


print('\nUsing next on iter(dataset) - Problems start here.')
dataset = new_dataset()
fill_dataset()
it = iter(dataset)

# Repeat the test-loop as behaviour strangely changes periodically
for repeat in range(5):

for i in range(30): # Flush count guess
s = next(it)
k, p = s.info.key.numpy().ravel(), s.info.priority.numpy().ravel()

# Iteratively halve the priorities
new_priorities = dict(zip(k, p * 0.999))
client.mutate_priorities(table.name, new_priorities)

s = next(it)
new_p = s.info.priority.numpy().ravel()

# Priority mutations now sync extremely slowly
if not np.isclose(p, new_p).all():
print(f'Priorities updated at flush-step {i}')
break
else:
# No break in for loop : not reached
print('No errors!')
```
Output:
```
before mutation
[[-1. -1. -1. -1. -1. 0.]
[-1. -1. -1. -1. -1. 0.]
[-1. -1. -1. -1. -1. 0.]]
sample priority: [1. 1. 1.]

after mutation
[[-1. -1. -1. -1. -1. 0.]
[-1. -1. -1. -1. -1. 0.]
[-1. -1. -1. -1. -1. 0.]]
sample priority: [0.5 0.5 0.5]

Using dataset.take
No errors!
No errors!
No errors!
No errors!
No errors!

Using next on iter(dataset) - Problems start here.
Priorities updated at flush-step 24
Priorities updated at flush-step 5
Priorities updated at flush-step 18
Priorities updated at flush-step 5
Priorities updated at flush-step 18
```

## Proposed Solution

The problem is immediately solved if `iter(dataset)` is called at each call to `next`. Because of this, I wasn't sure whether to post this issue here or in the tensorflow github, since the problem is with tf.data.Dataset. Personally I would suggest creating a wrapper around tf.data.Dataset for that either makes use of the `take` and `batch` API, or reinitialize the `iter` at every call. Because of how `reverb` implements sampling, reinitializing the dataset iterator should have no side-effects.

Example solution:
```python

print('\nReinitializing iter on every next call - Problem Solved.')
dataset = new_dataset()
fill_dataset()
it = iter(dataset) # Ignore this iterator

# Repeat the test-loop as behaviour strangely changes periodically
for repeat in range(5):

for i in range(30): # Flush count guess
s = next(iter(dataset)) # CHANGE: call iter(dataset) every time `next` is called
k, p = s.info.key.numpy().ravel(), s.info.priority.numpy().ravel()

# Iteratively halve the priorities
new_priorities = dict(zip(k, p * 0.999))
client.mutate_priorities(table.name, new_priorities)

s = next(iter(dataset)) # CHANGE: call iter(dataset) every time `next` is called
new_p = s.info.priority.numpy().ravel()

# Priority mutations now sync extremely slowly
if not np.isclose(p, new_p).all():
print(f'Priorities updated at flush-step {i}')
break
else:
# No break in for loop : not reached
print('No errors!')
```

Output: ( priorities are updated after every call, which is what we expected).
```
Reinitializing iter on every next call - Problem Solved.
Priorities updated at flush-step 0
Priorities updated at flush-step 0
Priorities updated at flush-step 0
Priorities updated at flush-step 0
Priorities updated at flush-step 0
```