feat: Mixed Experience Replay 🤝

[callumtilbury]

A simple utility to mix sampling of multiple buffers. Useful for offline-online stuff, and some off-policy variants that include portions of on-policy data (e.g. "combined experience replay," see here).

Important (& intentional) restrictions:

• The buffers must be of the same type. Why? We need to concatenate the samples, thus they need the same pytree structure. FlatBuffer returns an ExperiencePair, and this cannot be combined with a PrioritisedTrajectoryBufferSample, etc.
• We can ask for any ratio, [x,y,z,...], with a joint sample_batch_size when creating the mixer. We are still constrained by the underlying buffer sample functions, though. Suppose we have buffer_a which returns (4, ...) and buffer_b which returns (16, ...). We could create a mixer [1,1] with the sample_batch_size=6. In that case, we get 3 "batches" from buffer_a, and 3 batches from buffer_b. But if we ask for sample_batch_size=10 with ratio [1,1] (i.e. a batch size of 10/2 = 5 from each buffer), we'll only get the 4 batches from buffer_a, along with the 5 batches from buffer_b—so, a total sample_batch_size = 9. This is the idea of a "best effort"—we'll try grab enough batches of data, but only if possible. If not, we return a smaller batch than desired—as much as we can from each buffer.

It'd be great to test this out in a real system. Perhaps in Stoix, @EdanToledo? I can also look at stitching vaults together, etc.

See example notebook: https://colab.research.google.com/github/instadeepai/flashbax/blob/feat/mixed_experience_replay/examples/mixer_demonstration.ipynb (obviously won't run, unless you pip install the branch version, or run locally on a branch)

[eleninisioti]

Hello, I wanted to say that this is a great functionality. In my case I am interested in distributed RL with experience sharing among agents, so having a buffer where you can sample/add differently for personal use and sharing would be great. Is this a feature you are planning to merge soon? And if so, could you extend the notebook to show how you can add to the two buffers?

[EdanToledo]

@eleninisioti Hey, so yeah ideally we can merge this ASAP. I'm just waiting on one last thing from @callumtilbury but hes quite busy at the moment. If hes unable to complete it in the coming week, I'll take over and finish it. So hopefully sometime this week it will be merged :)

[EdanToledo]

looks good to me