Trajectory Writer performance worse than Legacy Writer

[Asad Jeewa]

We are working on Mava https://github.com/instadeepai/Mava which uses Reverb for adders (similar to the way Acme does). The trajectory writer seemed to be working slower than the legacy writer so I decided to do some benchmarking with dummy data and compare the writers directly, (isolated from Mava) by extending the Reverb tutorial code (https://gist.github.com/AsadJeewa/18ed6df8875a33c2cb7893738211a82d). It clearly shows that regardless of the number of episodes, sequence length, or other parameters, the trajectory writer is slower (as per my comments here (#78).
I have attached an example.

The first graph shows cumulative time for write() calls.

The second shows the total execution time.

I have more results to share as required.

[Sabela]

Hey Asad,

The trajectory writer that some more work in Python so it may cause a slowdown when using small steps (not sure if that is your case, though). One thing you can try is the new StructuredWriter, we found some performance improvements with that one and it should be possible to add it to your benchmark.

Some examples on how that writer is used are in the test (https://github.com/deepmind/reverb/blob/master/reverb/structured_writer_test.py).

Let us know if that helps!

Sabela.

[Asad Jeewa]

Thanks, I have been trying to add it to my benchmarks. Am I understanding this correctly? The point of the structured writer is to specify what to store via configs/ patterns. We can implement any writer that used the trajectory writers with the structured writer (and vice versa)

[Albin Cassirer]

It is possible to get the same behaviour using the StructuredWriter most of the time. If you have to do some custom logic at runtime to decide which part of an episode to add, or something along these lines, then it might not be possible.

[qstanczyk]

I think there were some improvements to the performance of Trajectory Writer, so closing this issue. Please reopen if this is still problematic.