import collections
import numpy as np

Transition = collections.namedtuple('Experience', field_names=['state', 'action', 'next_state', 'reward', 'is_game_on'])

class ReplayBuffer:
    def __init__(self, capacity):
        """回放缓冲区初始化"""
        self.buffer = collections.deque(maxlen=capacity)

    def push(self, transition):
        """存储经验"""
        self.buffer.append(transition)

    def sample(self, batch_size):
        """采样经验"""
        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        states, actions, next_states, rewards, dones = zip(*[self.buffer[idx] for idx in indices])
        return np.array(states), np.array(actions), np.array(next_states), np.array(rewards), np.array(dones)

    def __len__(self):
        """缓冲区长度"""
        return len(self.buffer)