<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> import time <EOL> import theano <EOL> floatX = theano . config . floatX <EOL> class DataSet ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , width , height , rng , max_steps = <NUM_LIT:1000> , phi_length = <NUM_LIT:4> ) : <EOL> """<STR_LIT>""" <EOL> self . width = width <EOL> self . height = height <EOL> self . max_steps = max_steps <EOL> self . phi_length = phi_length <EOL> self . rng = rng <EOL> self . imgs = np . zeros ( ( max_steps , height , width ) , dtype = '<STR_LIT>' ) <EOL> self . actions = np . zeros ( max_steps , dtype = '<STR_LIT>' ) <EOL> self . rewards = np . zeros ( max_steps , dtype = floatX ) <EOL> self . terminal = np . zeros ( max_steps , dtype = '<STR_LIT:bool>' ) <EOL> self . bottom = <NUM_LIT:0> <EOL> self . top = <NUM_LIT:0> <EOL> self . size = <NUM_LIT:0> <EOL> def add_sample ( self , img , action , reward , terminal ) : <EOL> """<STR_LIT>""" <EOL> self . imgs [ self . top ] = img <EOL> self . actions [ self . top ] = action <EOL> self . rewards [ self . top ] = reward <EOL> self . terminal [ self . top ] = terminal <EOL> if self . size == self . max_steps : <EOL> self . bottom = ( self . bottom + <NUM_LIT:1> ) % self . max_steps <EOL> else : <EOL> self . size += <NUM_LIT:1> <EOL> self . top = ( self . top + <NUM_LIT:1> ) % self . max_steps <EOL> def __len__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return max ( <NUM_LIT:0> , self . size - self . phi_length ) <EOL> def last_phi ( self ) : <EOL> """<STR_LIT>""" <EOL> indexes = np . arange ( self . top - self . phi_length , self . top ) <EOL> return self . imgs . take ( indexes , axis = <NUM_LIT:0> , mode = '<STR_LIT>' ) <EOL> def phi ( self , img ) : <EOL> """<STR_LIT>""" <EOL> indexes = np . arange ( self . top - self . phi_length + <NUM_LIT:1> , self . top ) <EOL> phi = np . empty ( ( self . phi_length , self . height , self . width ) , dtype = floatX ) <EOL> phi [ <NUM_LIT:0> : self . phi_length - <NUM_LIT:1> ] = self . imgs . take ( indexes , <EOL> axis = <NUM_LIT:0> , <EOL> mode = '<STR_LIT>' ) <EOL> phi [ - <NUM_LIT:1> ] = img <EOL> return phi <EOL> def random_batch ( self , batch_size ) : <EOL> """<STR_LIT>""" <EOL> states = np . zeros ( ( batch_size , <EOL> self . phi_length , <EOL> self . height , <EOL> self . width ) , <EOL> dtype = '<STR_LIT>' ) <EOL> actions = np . zeros ( ( batch_size , <NUM_LIT:1> ) , dtype = '<STR_LIT>' ) <EOL> rewards = np . zeros ( ( batch_size , <NUM_LIT:1> ) , dtype = floatX ) <EOL> terminal = np . zeros ( ( batch_size , <NUM_LIT:1> ) , dtype = '<STR_LIT:bool>' ) <EOL> next_states = np . zeros ( ( batch_size , <EOL> self . phi_length , <EOL> self . height , <EOL> self . width ) , <EOL> dtype = '<STR_LIT>' ) <EOL> count = <NUM_LIT:0> <EOL> while count < batch_size : <EOL> index = self . rng . randint ( self . bottom , <EOL> self . bottom + self . size - self . phi_length ) <EOL> initial_indices = np . arange ( index , index + self . phi_length ) <EOL> transition_indices = initial_indices + <NUM_LIT:1> <EOL> end_index = index + self . phi_length - <NUM_LIT:1> <EOL> if np . any ( self . terminal . take ( initial_indices [ <NUM_LIT:0> : - <NUM_LIT:1> ] , mode = '<STR_LIT>' ) ) : <EOL> continue <EOL> states [ count ] = self . imgs . take ( initial_indices , axis = <NUM_LIT:0> , mode = '<STR_LIT>' ) <EOL> actions [ count ] = self . actions . take ( end_index , mode = '<STR_LIT>' ) <EOL> rewards [ count ] = self . rewards . take ( end_index , mode = '<STR_LIT>' ) <EOL> terminal [ count ] = self . terminal . take ( end_index , mode = '<STR_LIT>' ) <EOL> next_states [ count ] = self . imgs . take ( transition_indices , <EOL> axis = <NUM_LIT:0> , <EOL> mode = '<STR_LIT>' ) <EOL> count += <NUM_LIT:1> <EOL> return states , actions , rewards , next_states , terminal <EOL> def simple_tests ( ) : <EOL> np . random . seed ( <NUM_LIT> ) <EOL> dataset = DataSet ( width = <NUM_LIT:2> , height = <NUM_LIT:3> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT:6> , phi_length = <NUM_LIT:4> ) <EOL> for i in range ( <NUM_LIT:10> ) : <EOL> img = np . random . randint ( <NUM_LIT:0> , <NUM_LIT> , size = ( <NUM_LIT:3> , <NUM_LIT:2> ) ) <EOL> action = np . random . randint ( <NUM_LIT:16> ) <EOL> reward = np . random . random ( ) <EOL> terminal = False <EOL> if np . random . random ( ) < <NUM_LIT> : <EOL> terminal = True <EOL> print '<STR_LIT>' , img <EOL> dataset . add_sample ( img , action , reward , terminal ) <EOL> print "<STR_LIT:I>" , dataset . imgs <EOL> print "<STR_LIT:A>" , dataset . actions <EOL> print "<STR_LIT:R>" , dataset . rewards <EOL> print "<STR_LIT:T>" , dataset . terminal <EOL> print "<STR_LIT>" , dataset . size <EOL> print <EOL> print "<STR_LIT>" , dataset . last_phi ( ) <EOL> print <EOL> print '<STR_LIT>' , dataset . random_batch ( <NUM_LIT:2> ) <EOL> def speed_tests ( ) : <EOL> dataset = DataSet ( width = <NUM_LIT> , height = <NUM_LIT> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT> , phi_length = <NUM_LIT:4> ) <EOL> img = np . random . randint ( <NUM_LIT:0> , <NUM_LIT> , size = ( <NUM_LIT> , <NUM_LIT> ) ) <EOL> action = np . random . randint ( <NUM_LIT:16> ) <EOL> reward = np . random . random ( ) <EOL> start = time . time ( ) <EOL> for i in range ( <NUM_LIT> ) : <EOL> terminal = False <EOL> if np . random . random ( ) < <NUM_LIT> : <EOL> terminal = True <EOL> dataset . add_sample ( img , action , reward , terminal ) <EOL> print "<STR_LIT>" , <NUM_LIT> / ( time . time ( ) - start ) <EOL> start = time . time ( ) <EOL> for i in range ( <NUM_LIT:200> ) : <EOL> a = dataset . random_batch ( <NUM_LIT:32> ) <EOL> print "<STR_LIT>" , <NUM_LIT:200> / ( time . time ( ) - start ) <EOL> print dataset . last_phi ( ) <EOL> def trivial_tests ( ) : <EOL> dataset = DataSet ( width = <NUM_LIT:2> , height = <NUM_LIT:1> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT:3> , phi_length = <NUM_LIT:2> ) <EOL> img1 = np . array ( [ [ <NUM_LIT:1> , <NUM_LIT:1> ] ] , dtype = '<STR_LIT>' ) <EOL> img2 = np . array ( [ [ <NUM_LIT:2> , <NUM_LIT:2> ] ] , dtype = '<STR_LIT>' ) <EOL> img3 = np . array ( [ [ <NUM_LIT:3> , <NUM_LIT:3> ] ] , dtype = '<STR_LIT>' ) <EOL> dataset . add_sample ( img1 , <NUM_LIT:1> , <NUM_LIT:1> , False ) <EOL> dataset . add_sample ( img2 , <NUM_LIT:2> , <NUM_LIT:2> , False ) <EOL> dataset . add_sample ( img3 , <NUM_LIT:2> , <NUM_LIT:2> , True ) <EOL> print "<STR_LIT>" , dataset . last_phi ( ) <EOL> print "<STR_LIT>" , dataset . random_batch ( <NUM_LIT:1> ) <EOL> def max_size_tests ( ) : <EOL> dataset1 = DataSet ( width = <NUM_LIT:3> , height = <NUM_LIT:4> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT:10> , phi_length = <NUM_LIT:4> ) <EOL> dataset2 = DataSet ( width = <NUM_LIT:3> , height = <NUM_LIT:4> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT:1000> , phi_length = <NUM_LIT:4> ) <EOL> for i in range ( <NUM_LIT:100> ) : <EOL> img = np . random . randint ( <NUM_LIT:0> , <NUM_LIT> , size = ( <NUM_LIT:4> , <NUM_LIT:3> ) ) <EOL> action = np . random . randint ( <NUM_LIT:16> ) <EOL> reward = np . random . random ( ) <EOL> terminal = False <EOL> if np . random . random ( ) < <NUM_LIT> : <EOL> terminal = True <EOL> dataset1 . add_sample ( img , action , reward , terminal ) <EOL> dataset2 . add_sample ( img , action , reward , terminal ) <EOL> np . testing . assert_array_almost_equal ( dataset1 . last_phi ( ) , <EOL> dataset2 . last_phi ( ) ) <EOL> print "<STR_LIT>" <EOL> def test_memory_usage_ok ( ) : <EOL> import memory_profiler <EOL> dataset = DataSet ( width = <NUM_LIT> , height = <NUM_LIT> , <EOL> rng = np . random . RandomState ( <NUM_LIT> ) , <EOL> max_steps = <NUM_LIT> , phi_length = <NUM_LIT:4> ) <EOL> last = time . time ( ) <EOL> for i in xrange ( <NUM_LIT> ) : <EOL> if ( i % <NUM_LIT> ) == <NUM_LIT:0> : <EOL> print i <EOL> dataset . add_sample ( np . random . random ( ( <NUM_LIT> , <NUM_LIT> ) ) , <NUM_LIT:1> , <NUM_LIT:1> , False ) <EOL> if i > <NUM_LIT> : <EOL> states , actions , rewards , next_states , terminals = dataset . random_batch ( <NUM_LIT:32> ) <EOL> if ( i % <NUM_LIT> ) == <NUM_LIT:0> : <EOL> print time . time ( ) - last <EOL> mem_usage = memory_profiler . memory_usage ( - <NUM_LIT:1> ) <EOL> print len ( dataset ) , mem_usage <EOL> last = time . time ( ) <EOL> def main ( ) : <EOL> speed_tests ( ) <EOL> test_memory_usage_ok ( ) <EOL> max_size_tests ( ) <EOL> simple_tests ( ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> main ( ) </s>
