import random from collections import namedtuple Transaction = namedtuple('Transaction', ('state', 'action', 'reward', 'next_state')) class Memory(object): def __init__(self, capacity): self.capacity = capacity self.memory = [] self.index = 0 def push(self, state, action, reward, next_state): if len(self.memory) < self.capacity: self.memory.append(None) self.memory[self.index] = (state, action, reward, next_state) self.index = (self.index + 1) % self.capacity def sample(self, batch_size): assert len(self.memory) >= self.capacity return random.sample(self.memory, batch_size) def __len__(self): return len(self.memory)