25 lines
730 B
Python
25 lines
730 B
Python
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)
|