utils.py

import numpy as np
import os

# Code based on:
# https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py


# Simple replay buffer
class ReplayBuffer(object):
    def __init__(self):
        self.storage = []

    # Expects tuples of (state, next_state, action, reward, done)
    def add(self, data):
        self.storage.append(data)

    def sample(self, batch_size):
        ind = np.random.randint(0, len(self.storage), size=batch_size)
        state, next_state, action, reward, done = [], [], [], [], []

        for i in ind:
            s, s2, a, r, d = self.storage[i]
            state.append(np.array(s, copy=False))
            next_state.append(np.array(s2, copy=False))
            action.append(np.array(a, copy=False))
            reward.append(np.array(r, copy=False))
            done.append(np.array(d, copy=False))

        return (
            np.array(state),
            np.array(next_state),
            np.array(action),
            np.array(reward).reshape(-1, 1),
            np.array(done).reshape(-1, 1)
        )

    def save(self, filename, directory):
        np.save(
            os.path.join(directory, '{}.npy'.format(filename)), self.storage
        )

    def load(self, filename, directory):
        self.storage = np.load(
            os.path.join(directory, '{}.npy'.format(filename))
        )