Add samples per second logging for `reverb_dataset.py`

[hartikainen]

trafficstars

Adds a more informative samples per second logging for reverb_dataset.py. Before:

Iteration 0 finished in 0.8698790073394775s
Iteration 1 finished in 0.9794516563415527s
Iteration 2 finished in 0.9291434288024902s

After:

Iteration 0 finished in 1.33s with 192_907.36 samples/s.
Iteration 1 finished in 1.37s with 187_186.11 samples/s.
Iteration 2 finished in 1.69s with 151_599.91 samples/s.