<reponame>Aetf/fathom<filename>fathom/deepq/database.py
from __future__ import absolute_import, print_function, division

import numpy as np


class database(object):
    def __init__(self, params):
        self.size = params['db_size']
        self.img_scale = params['img_scale']
        self.states = np.zeros([self.size, 84, 84], dtype='uint8')  # image dimensions
        self.actions = np.zeros(self.size, dtype='float32')
        self.terminals = np.zeros(self.size, dtype='float32')
        self.rewards = np.zeros(self.size, dtype='float32')
        self.bat_size = params['batch']
        self.bat_s = np.zeros([self.bat_size, 84, 84, 4])
        self.bat_a = np.zeros([self.bat_size])
        self.bat_t = np.zeros([self.bat_size])
        self.bat_n = np.zeros([self.bat_size, 84, 84, 4])
        self.bat_r = np.zeros([self.bat_size])

        self.counter = 0  # keep track of next empty state
        self.flag = False
        return

    def get_batches(self):
        for i in range(self.bat_size):
            idx = 0
            while idx < 3 or (idx > self.counter-2 and idx < self.counter+3):
                idx = np.random.randint(3, self.get_size()-1)
            self.bat_s[i] = np.transpose(self.states[idx-3:idx+1, :, :], (1, 2, 0))/self.img_scale
            self.bat_n[i] = np.transpose(self.states[idx-2:idx+2, :, :], (1, 2, 0))/self.img_scale
            self.bat_a[i] = self.actions[idx]
            self.bat_t[i] = self.terminals[idx]
            self.bat_r[i] = self.rewards[idx]
        #self.bat_s[0] = np.transpose(self.states[10:14,:,:],(1,2,0))/self.img_scale
        #self.bat_n[0] = np.transpose(self.states[11:15,:,:],(1,2,0))/self.img_scale
        #self.bat_a[0] = self.actions[13]
        #self.bat_t[0] = self.terminals[13]
        #self.bat_r[0] = self.rewards[13]

        return self.bat_s, self.bat_a, self.bat_t, self.bat_n, self.bat_r

    def insert(self, prevstate_proc, reward, action, terminal):
        self.states[self.counter] = prevstate_proc
        self.rewards[self.counter] = reward
        self.actions[self.counter] = action
        self.terminals[self.counter] = terminal
        # update counter
        self.counter += 1
        if self.counter >= self.size:
            self.flag = True
            self.counter = 0
        return

    def get_size(self):
        if self.flag == False:
            return self.counter
        else:
            return self.size

from __future__ import absolute_import, print_function, division

import numpy as np


class database(object):
    def __init__(self, params):
        self.size = params['db_size']
        self.img_scale = params['img_scale']
        self.states = np.zeros([self.size, 84, 84], dtype='uint8')  # image dimensions
        self.actions = np.zeros(self.size, dtype='float32')
        self.terminals = np.zeros(self.size, dtype='float32')
        self.rewards = np.zeros(self.size, dtype='float32')
        self.bat_size = params['batch']
        self.bat_s = np.zeros([self.bat_size, 84, 84, 4])
        self.bat_a = np.zeros([self.bat_size])
        self.bat_t = np.zeros([self.bat_size])
        self.bat_n = np.zeros([self.bat_size, 84, 84, 4])
        self.bat_r = np.zeros([self.bat_size])

        self.counter = 0  # keep track of next empty state
        self.flag = False
        return

    def get_batches(self):
        for i in range(self.bat_size):
            idx = 0
            while idx < 3 or (idx > self.counter-2 and idx < self.counter+3):
                idx = np.random.randint(3, self.get_size()-1)
            self.bat_s[i] = np.transpose(self.states[idx-3:idx+1, :, :], (1, 2, 0))/self.img_scale
            self.bat_n[i] = np.transpose(self.states[idx-2:idx+2, :, :], (1, 2, 0))/self.img_scale
            self.bat_a[i] = self.actions[idx]
            self.bat_t[i] = self.terminals[idx]
            self.bat_r[i] = self.rewards[idx]
        #self.bat_s[0] = np.transpose(self.states[10:14,:,:],(1,2,0))/self.img_scale
        #self.bat_n[0] = np.transpose(self.states[11:15,:,:],(1,2,0))/self.img_scale
        #self.bat_a[0] = self.actions[13]
        #self.bat_t[0] = self.terminals[13]
        #self.bat_r[0] = self.rewards[13]

        return self.bat_s, self.bat_a, self.bat_t, self.bat_n, self.bat_r

    def insert(self, prevstate_proc, reward, action, terminal):
        self.states[self.counter] = prevstate_proc
        self.rewards[self.counter] = reward
        self.actions[self.counter] = action
        self.terminals[self.counter] = terminal
        # update counter
        self.counter += 1
        if self.counter >= self.size:
            self.flag = True
            self.counter = 0
        return

    def get_size(self):
        if self.flag == False:
            return self.counter
        else:
            return self.size
# image dimensions # keep track of next empty state #self.bat_s[0] = np.transpose(self.states[10:14,:,:],(1,2,0))/self.img_scale #self.bat_n[0] = np.transpose(self.states[11:15,:,:],(1,2,0))/self.img_scale #self.bat_a[0] = self.actions[13] #self.bat_t[0] = self.terminals[13] #self.bat_r[0] = self.rewards[13] # update counter