<reponame>bbueno5000/SelfDrivingCarDemo<filename>self_driving_car_demo/self_driving_car_demo.py
"""
TODO: docstring
"""
import csv
import glob
import keras.callbacks
import keras.layers.core
import keras.layers.recurrent
import keras.models
import keras.optimizer_v2
import math
import matplotlib.pyplot as pyplot
import numpy
import os.path
import pygame
import pymunk
import pymunk.pygame_util
import random
import timeit

class GameState:
    """
    TODO: docstring
    """
    def __init__(self):
        """
        TODO: docstring
        """
        self.width = 1000
        self.height = 700
        pygame.init()
        self.screen = pygame.display.set_mode((self.width, self.height))
        self.clock = pygame.time.Clock()
        self.screen.set_alpha(None)
        self.show_sensors = True
        self.draw_screen = True
        self.crashed = False
        self.space = pymunk.Space()
        self.space.gravity = pymunk.Vec2d(0.0, 0.0)
        self.create_car(100, 100, 0.5)
        self.num_steps = 0
        static = [
            pymunk.Segment(self.space.static_body, (0, 1), (0, self.height), 1),
            pymunk.Segment(self.space.static_body, (1, self.height), (self.width, self.height), 1),
            pymunk.Segment(self.space.static_body, (self.width-1, self.height), (self.width-1, 1), 1),
            pymunk.Segment(self.space.static_body, (1, 1), (self.width, 1), 1)]
        for s in static:
            s.friction = 1.0
            s.group = 1
            s.collision_type = 1
            s.color = pygame.color.THECOLORS['red']
        self.space.add(*static)
        self.obstacles = []
        self.obstacles.append(self.create_obstacle(200, 350, 100))
        self.obstacles.append(self.create_obstacle(700, 200, 125))
        self.obstacles.append(self.create_obstacle(600, 600, 35))
        self.create_cat()

    def car_is_crashed(self, readings):
        """
        TODO: docstring
        """
        if readings[0] == 1 or readings[1] == 1 or readings[2] == 1:
            return True
        return False

    def create_cat(self):
        """
        TODO: docstring
        """
        inertia = pymunk.moment_for_circle(1, 0, 14, (0, 0))
        self.cat_body = pymunk.Body(1, inertia)
        self.cat_body.position = 50, self.height - 100
        self.cat_shape = pymunk.Circle(self.cat_body, 30)
        self.cat_shape.color = pygame.color.THECOLORS["orange"]
        self.cat_shape.elasticity = 1.0
        self.cat_shape.angle = 0.5
        direction = pymunk.vec2d.Vec2d(1, 0).rotated(self.cat_body.angle)
        self.space.add(self.cat_body, self.cat_shape)

    def create_car(self, x, y, r):
        """
        TODO: docstring
        """
        inertia = pymunk.moment_for_circle(1, 0, 14, (0, 0))
        self.car_body = pymunk.Body(1, inertia)
        self.car_body.position = x, y
        self.car_shape = pymunk.Circle(self.car_body, 25)
        self.car_shape.color = pygame.color.THECOLORS["green"]
        self.car_shape.elasticity = 1.0
        self.car_body.angle = r
        driving_direction = pymunk.vec2d.Vec2d(1, 0).rotated(self.car_body.angle)
        self.car_body.apply_impulse_at_world_point(driving_direction, (0, 0))
        self.space.add(self.car_body, self.car_shape)

    def create_obstacle(self, x, y, r):
        """
        TODO: docstring
        """
        c_body = pymunk.Body(100.0, 100.0)
        c_shape = pymunk.Circle(c_body, r)
        c_shape.elasticity = 1.0
        c_body.position = x, y
        c_shape.color = pygame.color.THECOLORS["blue"]
        self.space.add(c_body, c_shape)
        return c_body

    def frame_step(self, action):
        """
        TODO: docstring
        """
        if action == 0:
            self.car_body.angle -= 0.2
        elif action == 1:
            self.car_body.angle += 0.2
        if self.num_steps % 100 == 0:
            self.move_obstacles()
        if self.num_steps % 5 == 0:
            self.move_cat()
        driving_direction = pymunk.vec2d.Vec2d(1, 0).rotated(self.car_body.angle)
        self.car_body.velocity = 100 * driving_direction
        self.screen.fill(pygame.color.THECOLORS["black"])
        options = pymunk.pygame_util.DrawOptions(self.screen)
        self.space.debug_draw(options)
        self.space.step(1.0 / 10)
        if self.draw_screen:
            pygame.display.flip()
        self.clock.tick()
        x, y = self.car_body.position
        readings = self.get_sonar_readings(x, y, self.car_body.angle)
        state = numpy.array([readings])
        if self.car_is_crashed(readings):
            self.crashed = True
            reward = -500
            self.recover_from_crash(driving_direction)
        else:
            reward = -5 + int(self.sum_readings(readings) / 10)
        self.num_steps += 1
        return reward, state

    def get_arm_distance(self, arm, x, y, angle, offset):
        """
        Used to count the distance.
        """
        i = 0
        for point in arm:
            i += 1
            rotated_p = self.get_rotated_point(x, y, point[0], point[1], angle + offset)
            if rotated_p[0] <= 0 or rotated_p[1] <= 0 \
                    or rotated_p[0] >= self.width or rotated_p[1] >= self.height:
                return i
            else:
                obs = self.screen.get_at(rotated_p)
                if self.get_track_or_not(obs) != 0:
                    return i
            if self.show_sensors:
                pygame.draw.circle(self.screen, (255, 255, 255), (rotated_p), 2)
        return i
  
    def get_rotated_point(self, x_1, y_1, x_2, y_2, radians):
        """
        Rotate x_2, y_2 around x_1, y_1 by angle.
        """
        x_change = (x_2 - x_1) * math.cos(radians) + (y_2 - y_1) * math.sin(radians)
        y_change = (y_1 - y_2) * math.cos(radians) - (x_1 - x_2) * math.sin(radians)
        new_x = x_change + x_1
        new_y = self.height - (y_change + y_1)
        return int(new_x), int(new_y)
   
    def get_sonar_readings(self, x, y, angle):
        """
        Instead of using a grid of boolean(ish) sensors, sonar readings
        simply return N "distance" readings, one for each sonar
        we're simulating. The distance is a count of the first non-zero
        reading starting at the object. For instance, if the fifth sensor
        in a sonar "arm" is non-zero, then that arm returns a distance of 5.
        """
        readings = []
        arm_left = self.make_sonar_arm(x, y)
        arm_middle = arm_left
        arm_right = arm_left
        readings.append(self.get_arm_distance(arm_left, x, y, angle, 0.75))
        readings.append(self.get_arm_distance(arm_middle, x, y, angle, 0))
        readings.append(self.get_arm_distance(arm_right, x, y, angle, -0.75))
        if self.show_sensors:
            pygame.display.update()
        return readings
    
    def get_track_or_not(self, reading):
        """
        TODO: docstring
        """
        if reading == pygame.color.THECOLORS['black']:
            return 0
        return 1

    def make_sonar_arm(self, x, y):
        """
        TODO: docstring
        """
        spread = 10
        distance = 20
        arm_points = []
        for i in range(1, 40):
            arm_points.append((distance + x + (spread * i), y))
        return arm_points

    def move_cat(self):
        """
        TODO: docstring
        """
        speed = random.randint(20, 200)
        self.cat_body.angle -= random.randint(-1, 1)
        direction = pymunk.vec2d.Vec2d(1, 0).rotated(self.cat_body.angle)
        self.cat_body.velocity = speed * direction
    
    def move_obstacles(self):
        """
        TODO: docstring
        """
        for obstacle in self.obstacles:
            speed = random.randint(1, 5)
            direction = pymunk.vec2d.Vec2d(1, 0).rotated(
                self.car_body.angle + random.randint(-2, 2))
            obstacle.velocity = speed * direction

    def recover_from_crash(self, driving_direction):
        """
        We hit something, so recover.
        """
        while self.crashed:
            self.car_body.velocity = -100 * driving_direction
            self.crashed = False
            for i in range(10):
                self.car_body.angle += 0.2
                self.screen.fill(pygame.color.THECOLORS["red"])
                options = pymunk.pygame_util.DrawOptions(self.screen)
                self.space.debug_draw(options)
                self.space.step(1.0 / 10)
                if self.draw_screen:
                    pygame.display.flip()
                self.clock.tick()

    def sum_readings(self, readings):
        """
        Sum the number of non-zero readings.
        """
        tot = 0
        for i in readings:
            tot += i
        return tot

class Learning:
    """
    TODO: docstring
    """
    def __init__(self):
        """
        TODO: docstring
        """
        self.gamma = 0.9
        self.num_input = 3
        self.tuning = False

    def __call__(self):
        """
        TODO: docstring
        """
        if self.tuning:
            param_list = []
            nn_params = [[164, 150], [256, 256], [512, 512], [1000, 1000]]
            batchSizes = [40, 100, 400]
            buffers = [10000, 50000]
            for nn_param in nn_params:
                for batchSize in batchSizes:
                    for buffer in buffers:
                        params = {
                            'batchSize': batchSize, 'buffer': buffer, 'nn': nn_param}
                        param_list.append(params)
            for param_set in param_list:
                self.launch_learn(param_set)
        else:
            nn_param = [164, 150]
            params = {'batchSize': 100, 'buffer': 50000, 'nn': nn_param}
            model = NN.neural_net(self.num_input, nn_param)
            self.train_net(model, params)

    def launch_learn(self, params):
        """
        TODO: docstring
        """
        filename = self.params_to_filename(params)
        print('Trying %s' % filename)
        if not os.path.isfile('results\\sonar-frames\\loss_data-' + filename + '.csv'):
            open('results\\sonar-frames\\loss_data-' + filename + '.csv', 'a').close()
            print('Starting test.')
            model = NN.neural_net(self.num_input, params['nn'])
            self.train_net(model, params)
        else:
            print('Already tested.')

    def log_results(self, filename, data_collect, loss_log):
        """
        Save the results to a file so we can graph it later.
        """
        with open('results\\sonar-frames\\learn_data-' + filename + '.csv', 'w') as data_dump:
            wr = csv.writer(data_dump)
            wr.writerows(data_collect)
        with open('results\\sonar-frames\\loss_data-' + filename + '.csv', 'w') as lf:
            wr = csv.writer(lf)
            for loss_item in loss_log:
                wr.writerow(loss_item)

    def params_to_filename(self, params):
        """
        TODO: docstring
        """
        return str(params['nn'][0]) + '-' + str(params['nn'][1]) + '-' + \
            str(params['batchSize']) + '-' + str(params['buffer'])

    def process_minibatch(self, minibatch, model):
        """
        This does the heavy lifting, aka, the training.
        """
        X_train, y_train = list(), list()
        for memory in minibatch:
            old_state_m, action_m, reward_m, new_state_m = memory
            old_qval = model.predict(old_state_m, batch_size=1)
            newQ = model.predict(new_state_m, batch_size=1)
            maxQ = numpy.max(newQ)
            y = numpy.zeros((1, 3))
            y[:] = old_qval[:]
            if reward_m != -500:
                update = (reward_m + (self.gamma * maxQ))
            else:
                update = reward_m
            y[0][action_m] = update
            X_train.append(old_state_m.reshape(self.num_input,))
            y_train.append(y.reshape(3,))
        X_train = numpy.array(X_train)
        y_train = numpy.array(y_train)
        return X_train, y_train

    def train_net(self, model, params):
        """
        TODO: docstring
        """
        filename = self.params_to_filename(params)
        observe = 1000 # number of frames to observe before training
        epsilon = 1
        train_frames = 1000000 # number of frames to play
        batchSize = params['batchSize']
        buffer = params['buffer']
        max_car_distance = 0
        car_distance = 0
        t = 0
        data_collect = []
        replay = []  # stores tuples of (S, A, R, S')
        loss_log = []
        game_state = GameState()
        _, state = game_state.frame_step((2))
        start_time = timeit.default_timer()
        while t < train_frames:
            t += 1
            car_distance += 1
            if random.random() < epsilon or t < observe:
                action = numpy.random.randint(0, 3)
            else:
                qval = model.predict(state, batch_size=1)
                action = (numpy.argmax(qval))  # best
            reward, new_state = game_state.frame_step(action)
            replay.append((state, action, reward, new_state))
            if t > observe:
                if len(replay) > buffer:
                    replay.pop(0)
                minibatch = random.sample(replay, batchSize)
                X_train, y_train = self.process_minibatch(minibatch, model)
                history = LossHistory()
                model.fit(
                    X_train, y_train, batch_size=batchSize,
                    epochs=1, verbose=0, callbacks=[history])
                loss_log.append(history.losses)
            state = new_state
            if epsilon > 0.1 and t > observe:
                epsilon -= (1/train_frames)
            if reward == -500:
                data_collect.append([t, car_distance])
                if car_distance > max_car_distance:
                    max_car_distance = car_distance
                tot_time = timeit.default_timer() - start_time
                fps = car_distance / tot_time
                print('Max: %d at %d\tepsilon %f\t(%d)\t%f fps' % \
                    (max_car_distance, t, epsilon, car_distance, fps))
                car_distance = 0
                start_time = timeit.default_timer()
            if t % 25000 == 0:
                model.save_weights(
                    'saved-models\\' + filename + '-' + str(t) + '.h5', overwrite=True)
                print('Saving model %s - %d' % (filename, t))
        log_results(filename, data_collect, loss_log)

class LossHistory(keras.callbacks.Callback):
    """
    TODO: docstring
    """
    def on_batch_end(self, batch, logs={}):
        """
        TODO: docstring
        """
        self.losses.append(logs.get('loss'))

    def on_train_begin(self, logs={}):
        """
        TODO: docstring
        """
        self.losses = []

class NN:
    """
    Design comes from: http://outlace.com/Reinforcement-Learning-Part-3/
    """
    def lstm_net(num_sensors, load=False):
        """
        TODO: docstring
        """
        model = keras.models.Sequential()
        model.add(keras.layers.recurrent.LSTM(output_dim=512, input_dim=num_sensors, return_sequences=True))
        model.add(keras.layers.core.Dropout(0.2))
        model.add(keras.layers.recurrent.LSTM(output_dim=512, input_dim=512, return_sequences=False))
        model.add(keras.layers.core.Dropout(0.2))
        model.add(keras.layers.core.Dense(output_dim=3, input_dim=512))
        model.add(keras.layers.core.Activation('linear'))
        model.compile(loss='mean_squared_error', optimizer='rmsprop')
        return model

    def neural_net(num_sensors, params, load=''):
        """
        TODO: docstring
        """
        model = keras.models.Sequential()
        model.add(keras.layers.core.Dense(
            params[0], kernel_initializer='lecun_uniform', input_shape=(num_sensors,)))
        model.add(keras.layers.core.Activation('relu'))
        model.add(keras.layers.core.Dropout(0.2))
        model.add(keras.layers.core.Dense(params[1], kernel_initializer='lecun_uniform'))
        model.add(keras.layers.core.Activation('relu'))
        model.add(keras.layers.core.Dropout(0.2))
        model.add(keras.layers.core.Dense(3, kernel_initializer='lecun_uniform'))
        model.add(keras.layers.core.Activation('linear'))
        rms = keras.optimizer_v2.rmsprop.RMSProp()
        model.compile(loss='mse', optimizer=rms)
        if load:
            model.load_weights(load)
        return model

class Playing:
    """
    Once a model is trained, use this to play it.
    """
    def __call__(self):
        """
        TODO: docstring
        """
        saved_model = 'saved-models\\164-150-100-50000-25000.h5'
        model = NN.Denseneural_net(3, [164, 150], saved_model)
        play(model)

    def play(self, model):
        """
        TODO: docstring
        """
        car_distance = 0
        game_state = GameState()
        _, state = game_state.frame_step((2))
        while True:
            car_distance += 1
            action = (numpy.argmax(model.predict(state, batch_size=1)))
            _, state = game_state.frame_step(action)
            if car_distance % 1000 == 0:
                print('Current distance: %d frames.' % car_distance)

class Plotting:
    """
    Take the data in the results folder and plot it.
    """
    def __call__(self):
        """
        TODO: docstring
        """
        for f in glob.glob('results\\sonar-frames\\learn*.csv'):
            self.plot_file(f, 'learn')
        for f in glob.glob('results\\sonar-frames\\loss*.csv'):
            self.plot_file(f, 'loss')

    def movingaverage(self, y, window_size):
        """
        Moving average function from:
        http://stackoverflow.com/questions/11352047/finding-moving-average-from-data-points-in-python
        """
        window = numpy.ones(int(window_size))/float(window_size)
        return numpy.convolve(y, window, 'same')

    def plot_file(self, filename, type='loss'):
        """
        DOCSTRING
        """
        with open(f, 'r') as csvfile:
            reader = csv.reader(csvfile)
            y = []
            for row in reader:
                if type == 'loss':
                    y.append(float(row[0]))
                else:
                    y.append(float(row[1]))
            if len(y) == 0:
                return
            print(self.readable_output(f))
            if type == 'loss':
                window = 100
            else:
                window = 10
            y_av = self.movingaverage(y, window)
            arr = numpy.array(y_av)
            if type == 'loss':
                print("%f\t%f\n" % (arr.min(), arr.mean()))
            else:
                print("%f\t%f\n" % (arr.max(), arr.mean()))
            pyplot.clf()
            pyplot.title(f)
            if type == 'loss':
                pyplot.plot(y_av[:-50])
                pyplot.ylabel('Smoothed Loss')
                pyplot.ylim(0, 5000)
                pyplot.xlim(0, 250000)
            else:
                pyplot.plot(y_av[:-5])
                pyplot.ylabel('Smoothed Distance')
                pyplot.ylim(0, 4000)
            pyplot.savefig(f + '.png', bbox_inches='tight')

    def readable_output(self, filename):
        """
        DOCSTRING
        """
        readable = ''
        f_parts = filename.split('-')
        if f_parts[0] == 'learn_data':
            readable += 'distance: '
        else:
            readable += 'loss: '
        readable += f_parts[1] + ', ' + f_parts[2] + ' | '
        readable += f_parts[3] + ' | '
        readable += f_parts[4].split('.')[0]
        return readable

if __name__ == '__main__':
    Learning()()
