aux_functions.py

# Author: Aqeel Anwar(ICSRL)
# Created: 10/14/2019, 12:50 PM
# Email: aqeel.anwar@gatech.edu
import numpy as np
import nvidia_smi
import os, subprocess, psutil
import math
import random
import time
import airsim
import pygame
from configs.read_cfg import read_cfg
import matplotlib.pyplot as plt
from PIL import Image
import cv2
from skimage.util import random_noise

def close_env(env_process):
    process = psutil.Process(env_process.pid)
    for proc in process.children(recursive=True):
        proc.kill()
    process.kill()


def save_network_path(cfg, algorithm_cfg):
    # Save the network to the directory network_path
    weights_type = 'Imagenet'
    bslash = '\\'
    if algorithm_cfg.custom_load == True:
        algorithm_cfg.network_path = 'models\\trained\\' + cfg.env_type + bslash + cfg.env_name + bslash + 'CustomLoad' + bslash + algorithm_cfg.train_type + bslash
    else:
        algorithm_cfg.network_path = 'models\\trained\\' + cfg.env_type + bslash + cfg.env_name + bslash + weights_type + bslash + algorithm_cfg.train_type + bslash

    if not os.path.exists(algorithm_cfg.network_path):
        os.makedirs(algorithm_cfg.network_path)

    return cfg, algorithm_cfg


def communicate_across_agents(agent, name_agent_list, algorithm_cfg):
    name_agent = name_agent_list[0]
    update_done = False
    if 'GlobalLearningGlobalUpdate' in algorithm_cfg.distributed_algo:
        # No need to do anything
        update_done = True

    elif algorithm_cfg.distributed_algo == 'LocalLearningGlobalUpdate':
        agent_on_same_network = name_agent_list
        agent[name_agent].network_model.initialize_graphs_with_average(agent, agent_on_same_network)

    elif algorithm_cfg.distributed_algo == 'LocalLearningLocalUpdate':
        agent_connectivity_graph = []
        for j in range(int(np.floor(len(name_agent_list) / algorithm_cfg.average_connectivity))):
            div1 = random.sample(name_agent_list, algorithm_cfg.average_connectivity)
            # print(div1)
            agent_connectivity_graph.append(div1)
            name_agent_list = list(set(name_agent_list) - set(div1))

        if name_agent_list:
            agent_connectivity_graph.append(name_agent_list)


        for agent_network in agent_connectivity_graph:
            agent_on_same_network = agent_network
            agent[name_agent].network_model.initialize_graphs_with_average(agent, agent_on_same_network)

    return update_done

def start_environment(env_name):
    print_orderly('Environment', 80)
    env_folder = os.path.dirname(os.path.abspath(__file__)) + "\\unreal_envs\\" + env_name + "\\"
    path = env_folder + env_name + ".exe"
    print(path)
    # env_process = []
    env_process = subprocess.Popen(path)
    time.sleep(5)
    print("Successfully loaded environment: " + env_name)

    return env_process, env_folder


def initialize_infer(env_cfg, client, env_folder):

    if not os.path.exists(env_folder+'results'):
        os.makedirs(env_folder+'results')

    # Mapping floor to 0 height
    f_z = env_cfg.floor_z/100
    c_z = (env_cfg.ceiling_z-env_cfg.floor_z)/100
    p_z = (env_cfg.player_start_z-env_cfg.floor_z)/100

    plt.ion()
    fig_z = plt.figure()
    ax_z = fig_z.add_subplot(111)
    line_z, = ax_z.plot(0, 0)
    ax_z.set_ylim(0, c_z)
    plt.title("Altitude variation")

    # start_posit = client.simGetVehiclePose()

    fig_nav = plt.figure()
    ax_nav = fig_nav.add_subplot(111)
    img = plt.imread(env_folder+ env_cfg.floorplan)
    ax_nav.imshow(img)
    plt.axis('off')
    plt.title("Navigational map")
    plt.plot(env_cfg.o_x, env_cfg.o_y, 'b*', linewidth=20)
    nav, = ax_nav.plot(env_cfg.o_x, env_cfg.o_y)


    return p_z,f_z, fig_z, ax_z, line_z, fig_nav, ax_nav, nav

def translate_action(action, num_actions):
    # action_word = ['Forward', 'Right', 'Left', 'Sharp Right', 'Sharp Left']
    sqrt_num_actions = np.sqrt(num_actions)
    # ind = np.arange(sqrt_num_actions)
    if sqrt_num_actions % 2 == 0:
        v_string = list('U' * int((sqrt_num_actions - 1) / 2) + 'D' * int((sqrt_num_actions - 1) / 2))
        h_string = list('L' * int((sqrt_num_actions - 1) / 2) + 'R' * int((sqrt_num_actions - 1) / 2))
    else:
        v_string = list('U' * int(sqrt_num_actions / 2) + 'F' + 'D' * int(sqrt_num_actions / 2))
        h_string = list('L' * int(sqrt_num_actions / 2) + 'F' + 'R' * int(sqrt_num_actions / 2))

    v_ind = int(action[0]/sqrt_num_actions)
    h_ind = int(action[0]%sqrt_num_actions)
    action_word = v_string[v_ind] + str(int(np.ceil(abs((sqrt_num_actions-1)/2-v_ind)))) + '-' + h_string[h_ind]+str(int(np.ceil(abs((sqrt_num_actions-1)/2-h_ind))))

    return action_word

def get_errors(data_tuple, choose, ReplayMemory, input_size, agent, target_agent, gamma, Q_clip):

    _, Q_target, _, err, _ = minibatch_double(data_tuple, len(data_tuple), choose,  ReplayMemory, input_size, agent, target_agent, gamma, Q_clip)

    return err


def minibatch_double(data_tuple, batch_size, choose, ReplayMemory, input_size, agent, target_agent, gamma, Q_clip):
    # Needs NOT to be in DeepAgent
    # NO TD error term, and using huber loss instead
    # Bellman Optimality equation update, with less computation, updated

    if batch_size==1:
        train_batch = data_tuple
        idx=None
    else:
        batch = ReplayMemory.sample(batch_size)
        train_batch = np.array([b[1][0] for b in batch])
        idx = [b[0] for b in batch]


    actions = np.zeros(shape=(batch_size), dtype=int)
    crashes = np.zeros(shape=(batch_size))
    rewards = np.zeros(shape=batch_size)
    curr_states = np.zeros(shape=(batch_size, input_size, input_size, 3))
    new_states = np.zeros(shape=(batch_size, input_size, input_size, 3))
    for ii, m in enumerate(train_batch):
        curr_state_m, action_m, new_state_m, reward_m, crash_m = m
        curr_states[ii, :, :, :] = curr_state_m[...]
        actions[ii] = action_m
        new_states[ii,:,:,:] = new_state_m
        rewards[ii] = reward_m
        crashes[ii] = crash_m

    #
    # oldQval = np.zeros(shape = [batch_size, num_actions])
    if choose:
        oldQval_A = target_agent.network_model.Q_val(curr_states)
        newQval_A = target_agent.network_model.Q_val(new_states)
        newQval_B = agent.network_model.Q_val(new_states)
    else:
        oldQval_A = agent.network_model.Q_val(curr_states)
        newQval_A = agent.network_model.Q_val(new_states)
        newQval_B = target_agent.network_model.Q_val(new_states)


    TD = np.zeros(shape=[batch_size])
    err = np.zeros(shape=[batch_size])
    Q_target = np.zeros(shape=[batch_size])

    term_ind = np.where(rewards==-1)[0]
    nonterm_ind = np.where(rewards!=-1)[0]

    TD[nonterm_ind] = rewards[nonterm_ind] + gamma* newQval_B[nonterm_ind, np.argmax(newQval_A[nonterm_ind], axis=1)] - oldQval_A[nonterm_ind, actions[nonterm_ind].astype(int)]
    TD[term_ind] = rewards[term_ind]

    if Q_clip:
        TD_clip = np.clip(TD, -1, 1)
    else:
        TD_clip = TD

    Q_target[nonterm_ind] = oldQval_A[nonterm_ind, actions[nonterm_ind].astype(int)] + TD_clip[nonterm_ind]
    Q_target[term_ind] = TD_clip[term_ind]

    err=abs(TD) # or abs(TD_clip)
    return curr_states, Q_target, actions, err, idx

def policy_REINFORCE(curr_state, agent):
    action = agent.network_model.action_selection(curr_state)
    action_type = 'Prob'
    return action[0], action_type

def train_REINFORCE(data_tuple, batch_size, agent, lr, input_size, gamma, epi_num):
    episode_len = len(data_tuple)

    curr_states = np.zeros(shape=(episode_len, input_size, input_size, 3))
    actions = np.zeros(shape=(episode_len), dtype=int)
    crashes = np.zeros(shape=(episode_len))
    rewards = np.zeros(shape=episode_len)

    for ii, m in enumerate(data_tuple):
        curr_state_m, action_m, reward_m, crash_m = m
        curr_states[ii, :, :, :] = curr_state_m[...]
        actions[ii] = action_m
        rewards[ii] = reward_m
        crashes[ii] = crash_m

    Gs = np.zeros(episode_len)
    r = 0
    for episode_step in range(episode_len - 1, -1, -1):
        r = rewards[episode_step] + r * gamma
        Gs[episode_step] = r

    # Normalize the reward to reduce variance in training
    Gs -= np.mean(Gs)
    Gs /= (np.std(Gs) + 1e-8)

    num_batches = int(np.ceil(episode_len / batch_size))
    for i in range(num_batches):
        if i != num_batches - 1:
            x = curr_states[i * batch_size:(i + 1) * batch_size, :, :, :]
            G = Gs[i * batch_size:(i + 1) * batch_size]
            action = actions[i * batch_size:(i + 1) * batch_size]
        else:
            x = curr_states[i * batch_size:, :, :, :]
            G = Gs[i * batch_size:]
            action = actions[i * batch_size:]

        G = np.array([G])
        G = G.T

        # Restructure array
        action = np.array([action])
        action = action.T

        # Get the baseline value
        B = agent.network_model.get_baseline(x)
        # Train the baseline network
        agent.network_model.train_baseline(x, G, action, lr, epi_num)
        # Train policy network
        agent.network_model.train_policy(x, action, B, G, lr, epi_num)


def policy(epsilon, curr_state, iter, b, epsilon_model, wait_before_train, num_actions, agent):
    qvals=[]

    epsilon_ceil=0.95
    if epsilon_model=='linear':
        epsilon = epsilon_ceil* (iter-wait_before_train) / (b-wait_before_train)
        if epsilon > epsilon_ceil:
            epsilon = epsilon_ceil

    elif epsilon_model=='exponential':
        epsilon = 1- math.exp(-2/(b-wait_before_train) * (iter-wait_before_train) )
        if epsilon > epsilon_ceil:
            epsilon = epsilon_ceil

    if random.random() > epsilon:
        sss =curr_state.shape
        action = np.random.randint(0, num_actions, size = sss[0], dtype=np.int32)
        action_type = 'Rand'
    else:
        # Use NN to predict action
        action = agent.network_model.action_selection(curr_state)
        action_type = 'Pred'
        # print(action_array/(np.mean(action_array)))
    return action, action_type, epsilon, qvals

def reset_to_initial(level, reset_array, client, vehicle_name):
    # client.moveByVelocityAsync(vx=0, vy=0, vz=0, duration=0.01, vehicle_name=vehicle_name)
    reset_pos = reset_array[vehicle_name][level]
    #print(reset_pos, level)
    # reset_pos = p

    client.simSetVehiclePose(reset_pos, ignore_collison=True, vehicle_name=vehicle_name)
    time.sleep(0.05)

def print_orderly(str, n):
    print('')
    hyphens = '-' * int((n - len(str)) / 2)
    print(hyphens + ' ' + str + ' ' + hyphens)

def connect_drone(ip_address='127.0.0.0', phase='infer', num_agents=1, client=[]):
    if client != []:
        client.reset()
    print_orderly('Drone', 80)
    client = airsim.MultirotorClient(ip=ip_address, timeout_value=10)
    client.confirmConnection()
    time.sleep(1)

    old_posit = {}
    for agents in range(num_agents):
        name_agent = "drone" + str(agents)
        client.enableApiControl(True, name_agent)
        client.armDisarm(True, name_agent)
        # time.sleep(1)
        client.takeoffAsync(vehicle_name=name_agent)
        time.sleep(1)
        old_posit[name_agent] = client.simGetVehiclePose(vehicle_name=name_agent)
        #print_orderly(name_agent, 40)

    initZ = old_posit[name_agent].position.z_val

    return client, old_posit, initZ


def get_SystemStats(process, NVIDIA_GPU):
    if NVIDIA_GPU:
        deviceCount = nvidia_smi.nvmlDeviceGetCount()
        gpu_memory= []
        gpu_utilization = []
        for i in range(0, deviceCount):
            handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
            gpu_stat = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
            gpu_memory.append(gpu_stat.memory)
            gpu_utilization.append(gpu_stat.gpu)
    else:
        gpu_memory = []
        gpu_utilization = []

    sys_memory = process.memory_info()[0] / 2. ** 30

    return gpu_memory, gpu_utilization, sys_memory

def get_MonocularImageRGB(client, vehicle_name):

    responses1 = client.simGetImages([
        airsim.ImageRequest('front_center', airsim.ImageType.Scene, False,
                            False)], vehicle_name=vehicle_name)  # scene vision image in uncompressed RGBA array

    response = responses1[0]
    img1d = np.fromstring(response.image_data_uint8, dtype=np.uint8)  # get numpy array
    img_rgba = img1d.reshape(response.height, response.width, 3)
    img = Image.fromarray(img_rgba)
    img_rgb = img.convert('RGB')
    camera_image_rgb = np.asarray(img_rgb)
    camera_image = camera_image_rgb

    return camera_image

def get_StereoImageRGB(client, vehicle_name):
    camera_image=[]
    responses = client.simGetImages(
        [
        airsim.ImageRequest('front_left', airsim.ImageType.Scene, False, False),
        airsim.ImageRequest('front_right', airsim.ImageType.Scene, False, False)
        ], vehicle_name=vehicle_name)

    for i in range(2):
        response = responses[i]
        img1d = np.fromstring(response.image_data_uint8, dtype=np.uint8)  # get numpy array
        img_rgba = img1d.reshape(response.height, response.width, 3)
        img = Image.fromarray(img_rgba)
        img_rgb = img.convert('RGB')
        camera_image_rgb = np.asarray(img_rgb)
        camera_image.append(camera_image_rgb)

    return camera_image

def get_CustomImage(client, vehicle_name, camera_name):
    responses1 = client.simGetImages([
        airsim.ImageRequest(camera_name, airsim.ImageType.Scene, False,
                            False)], vehicle_name=vehicle_name)  # scene vision image in uncompressed RGBA array

    response = responses1[0]
    img1d = np.fromstring(response.image_data_uint8, dtype=np.uint8)  # get numpy array
    img_rgba = img1d.reshape(response.height, response.width, 3)
    img = Image.fromarray(img_rgba)
    img_rgb = img.convert('RGB')
    camera_image_rgb = np.asarray(img_rgb)
    camera_image = camera_image_rgb

    return camera_image


# def get_image(client, vehicle_name, camera_type, first_frame, last_frame):
#     responses1 = client.simGetImages([  # depth visualization image
#         airsim.ImageRequest("1", airsim.ImageType.Scene, False,
#                             False)], vehicle_name=vehicle_name)  # scene vision image in uncompressed RGBA array
#
#     response = responses1[0]
#     img1d = np.fromstring(response.image_data_uint8, dtype=np.uint8)  # get numpy array
#     img_rgba = img1d.reshape(response.height, response.width, 3)
#     img = Image.fromarray(img_rgba)
#     img_rgb = img.convert('RGB')
#     camera_image_rgb = np.asarray(img_rgb)
#
#     if camera_type == 'optical':
#         camera_image = camera_image_rgb
#
#     if camera_type == 'DVS':
#         # camera_image = cv2.normalize(camera_image_rgb, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
#         frame1 = cv2.cvtColor(camera_image_rgb, cv2.COLOR_BGR2GRAY)
#         # frame23 = cv2.normalize(frame1, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
#         frame = np.uint8(np.log1p(frame1))
#         frame = cv2.normalize(frame, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
#
#         if first_frame:
#             camera_image = frame
#             first_frame = False
#         else:
#             camera_image = frame - last_frame
#         # ret, thresh1 = cv2.threshold(display_frame, 0.2, 0.8, cv2.THRESH_BINARY)
#         # display_frame1 = cv2.bitwise_and(display_frame, thresh1)
#         last_frame = frame
#
#         camera_image = random_noise(camera_image, mode='s&p', amount=0.005)
#         camera_image = cv2.cvtColor(camera_image, cv2.COLOR_GRAY2BGR)
#
#         cv2.imshow('rgb', camera_image_rgb)
#         cv2.imshow('dvs', camera_image)
#         cc=1

    # return camera_image, first_frame, last_frame


def blit_text(surface, text, pos, font, color=pygame.Color('black')):
    words = [word.split(' ') for word in text.splitlines()]  # 2D array where each row is a list of words.
    space = font.size(' ')[0]  # The width of a space.
    max_width, max_height = surface.get_size()
    x, y = pos
    for line in words:
        for word in line:
            word_surface = font.render(word, 0, color)
            word_width, word_height = word_surface.get_size()
            if x + word_width >= max_width:
                x = pos[0]  # Reset the x.
                y += word_height  # Start on new row.
            surface.blit(word_surface, (x, y))
            x += word_width + space
        x = pos[0]  # Reset the x.
        y += word_height  # Start on new row.

def pygame_connect(phase):
    pygame.init()

    if phase == 'train':
        img_path = 'images\\train_keys.png'
    elif phase == 'infer':
        img_path = 'images\\infer_keys.png'
    img = pygame.image.load(img_path)

    screen = pygame.display.set_mode(img.get_rect().size)

    screen.blit(img, (0, 0))
    pygame.display.set_caption('DLwithTL')
    # screen.fill((21, 116, 163))
    # text = 'Supported Keys:\n'
    # font = pygame.font.SysFont('arial', 32)
    # blit_text(screen, text, (20, 20), font, color = (214, 169, 19))
    # pygame.display.update()
    #
    # font = pygame.font.SysFont('arial', 24)
    # text = 'R - Reconnect unreal\nbackspace - Pause/play\nL - Update configurations\nEnter - Save Network'
    # blit_text(screen, text, (20, 70), font, color=(214, 169, 19))
    pygame.display.update()

    return screen

def check_user_input(active, automate, agent, client, old_posit, initZ, fig_z, fig_nav, env_folder,cfg, algorithm_cfg):
    # algorithm_cfg.learning_rate, algorithm_cfg.epsilon,algorithm_cfg.network_path,cfg.mode,
    for event in pygame.event.get():

        if event.type == pygame.QUIT:
            active = False
            pygame.quit()

        # Training keys control
        if event.type == pygame.KEYDOWN and cfg.mode =='train':
            if event.key == pygame.K_l:
                # Load the parameters - epsilon
                path = 'configs\\' + cfg.algorithm + '.cfg'
                algorithm_cfg = read_cfg(config_filename=path, verbose=False)
                cfg, algorithm_cfg = save_network_path(cfg=cfg, algorithm_cfg=algorithm_cfg)
                print('Updated Parameters')

            if event.key == pygame.K_RETURN:
                # take_action(-1)
                automate = False
                print('Saving Model')
                # agent.save_network(iter, save_path, ' ')
                agent.network_model.save_network(algorithm_cfg.network_path, episode='user')
                # agent.save_data(iter, data_tuple, tuple_path)


            if event.key == pygame.K_BACKSPACE:
                automate = automate ^ True

            if event.key == pygame.K_r:
                client, old_posit, initZ = connect_drone(ip_address=cfg.ip_address, phase=cfg.mode,
                                                         num_agents=cfg.num_agents)

                agent.client = client


            # Set the routine for manual control if not automate
            if not automate:
                # print('manual')
                # action=[-1]
                if event.key == pygame.K_UP:
                    action = 0
                elif event.key == pygame.K_RIGHT:
                    action = 1
                elif event.key == pygame.K_LEFT:
                    action = 2
                elif event.key == pygame.K_d:
                    action = 3
                elif event.key == pygame.K_a:
                    action = 4
                elif event.key == pygame.K_DOWN:
                    action = -2
                elif event.key == pygame.K_y:
                    pos = client.getPosition()

                    client.moveToPosition(pos.x_val, pos.y_val, 3 * initZ, 1)
                    time.sleep(0.5)
                elif event.key == pygame.K_h:
                    client.reset()
                # agent.take_action(action)

        elif event.type == pygame.KEYDOWN and cfg.mode == 'infer':
            if event.key == pygame.K_s:
                # Save the figures
                file_path = env_folder + 'results\\'
                fig_z.savefig(file_path+'altitude_variation.png', dpi=1000)
                fig_nav.savefig(file_path+'navigation.png', dpi=1000)
                print('Figures saved')

            if event.key == pygame.K_BACKSPACE:
                client.moveByVelocityAsync(vx=0, vy=0, vz=0, duration=0.1)
                automate = automate ^ True

    return active, automate, algorithm_cfg, client