feat: vault

[callumtilbury]

What?

Vault is an efficient mechanism for saving flashbax buffers to persistent data storage.

How?

Vault uses tensorstore, which is also used by Google's orbax checkpointing library. Tensorstore has a useful ability to read and write slices of buffers at a time. i.e. I could have a 100TB file saved on disk, but easily and efficiently access the first element. Building on this feature, with flashbax buffers saved in the form of (Batch, Time, Experience), we slice data along the time axis.

Vault consumes flashbax buffer states and writes to disk from the last received time index to the buffer's current time index. This must be done before the ring buffer overwrites any stale data that has not yet been written to the vault. All other bookkeeping is done by vault itself, and using it is as simple as: v.write(buffer_state)

It is usually helpful to look at the demonstrative notebook, which adds timesteps to the buffer, writing to the vault after each one. The following output is yielded:

------------------
Buffer state:
[[[0.]
  [0.]
  [0.]
  [0.]
  [0.]]]

Vault state:
[]
------------------
------------------
Buffer state:
[[[1.]
  [0.]
  [0.]
  [0.]
  [0.]]]

Vault state:
[[[1.]]]
------------------
------------------
Buffer state:
[[[1.]
  [2.]
  [0.]
  [0.]
  [0.]]]

Vault state:
[[[1.]
  [2.]]]
------------------
------------------
Buffer state:
[[[1.]
  [2.]
  [3.]
  [0.]
  [0.]]]

Vault state:
[[[1.]
  [2.]
  [3.]]]
------------------
------------------
Buffer state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [0.]]]

Vault state:
[[[1.]
  [2.]
  [3.]
  [4.]]]
------------------
------------------
Buffer state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [5.]]]

Vault state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [5.]]]
------------------
------------------
Buffer state:
[[[6.]
  [2.]
  [3.]
  [4.]
  [5.]]]

Vault state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [5.]
  [6.]]]
------------------
------------------
Buffer state:
[[[6.]
  [7.]
  [3.]
  [4.]
  [5.]]]

Vault state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [5.]
  [6.]
  [7.]]]
------------------
------------------
Buffer state:
[[[6.]
  [7.]
  [8.]
  [4.]
  [5.]]]

Vault state:
[[[1.]
  [2.]
  [3.]
  [4.]
  [5.]
  [6.]
  [7.]
  [8.]]]
------------------

Why?

Variety of reasons—but mainly useful for Offline RL methods, and saving/loading buffers from checkpoints.

[sash-a]

Minor nitpicks and a suggestion for the list issue. Review still in progress though

[EdanToledo]

Looks great to me! thank you so much Callum for your hard work and dedicated time to support the offline RL community. They will be very appreciative and thankful.

Stale & now Sasha is on leave :')