pool.py

import random

class Pool():
    def __init__(self, size):
        self.size = size
        self.data = [None] * size
        self.idx  = 0

        self.sum_len = 0

        self.total = 0

    def put(self, x):
        if(self.total >= self.size):
            old_x = self.data[self.idx]
            self.sum_len -= len(old_x[0])

        self.sum_len += len(x[0])

        self.data[self.idx] = x
        self.idx = (self.idx + 1) % self.size
        self.total += 1

    ''' Sample a batch of #size episodes. '''
    def sample(self, size):
        return random.choices(self.data, k=size)

    ''' Samples a batch of episodes. Size of total steps is close to #size.'''
    def sample_steps(self, size):
        avg_len = self.sum_len / self.size
        eps_to_fetch = int(size / avg_len)

        return random.choices(self.data, k=eps_to_fetch)