Write custom serialisation for experience replay

Although it is now possible to save/load experience replay memories (https://github.com/Kaixhin/PlaNet/issues/3), naively using torch.save fails with large memories. Dealing with this would require custom serialisation code.

[Kaixhin]

I was just talking about https://github.com/Kaixhin/PlaNet/blob/master/memory.py#L15-L18 , but yes, anything that needs to be stored to recover the whole ExperienceReplay class.

[alec-tschantz]

Oh nice, this absolutely seems viable. (10**6, 3, 64, 64) is 12 GB on the hard-drive but reading and writing are mostly instant. As a minimal working example:

class Buffer(object):
    def __init__(self, buffer_size=10 ** 6):
        self.buffer_size = buffer_size
        self.obs_shape = (buffer_size, 3, 64, 64)
        self.obs_path = "obs.dat"
        self.idx = 0

    def create_new_file(self):
        obs_f = np.memmap(
            self.obs_path, dtype=np.uint8, mode="w+", shape=self.obs_shape
        )
        del obs_f

    def add(self, obs):
        obs_f = np.memmap(
            self.obs_path, dtype=np.uint8, mode="w+", shape=self.obs_shape
        )
        obs_f[self.idx] = obs
        del obs_f
        self.idx += 1

    def sample(self, idxs):
        obs_f = np.memmap(self.obs_path, dtype=np.uint8, mode="r", shape=self.obs_shape)
        data = obs_f[idxs]
        del obs_f
        return data

if __name__ == "__main__":
    buffer = Buffer()
    for i in range(10 ** 6):
        obs = (np.random.rand(3, 64, 64) * 255).astype(np.uint8)
        buffer.add(obs)
    idxs = np.random.randint(0, 10 ** 4, size=100)
    data = buffer.sample(idxs)

I’ll get round to integrating it with your buffer at some point - if you think it’d be worthwhile.