[IDEA] Easiest way to implement Hindsight Relabeling?

Hi all,

I was wondering if anyone had any thoughts about ways to implement Hindsight experience replay relabeling in to a replay buffer for Goal Conditioned implementations. In essence, it would look a lot like the trajectory buffer with a few key differences:

- We insert (in a batched format) trajectories from an agent's experience that are labeled with (s, a, s', r, g)
- (either at sample time or insertion time) we extract the goals achieved during the trajectory and populate the replay buffer with n copies of the trajectory, each with modified g values as if the agent "intended" to accomplish those goals
- The goal relabeling often requires the entire trajectory for context (e.g. you might need to look ahead to the future to know what to label the current state with).

It would be nice if you didn't have to relabel every time you sample, but instead re-insert relabeled trajectories into the replay buffer to be further sampled. Also, while we need the full trajectory for relabeling context, it might not be needed for actual learning algorithm (.e.g if we were doing a goal-conditioned DQN).

Does anyone have any thoughts about how to best go about implementing something like this using flashbax?

Thanks in advance!