general_test16.py

import sys
import argparse
import json
import matplotlib.pyplot as plt
import random
import numpy as np
import torch
from torch.autograd import Variable

from dataset.dataset import *
from utility.utils import *
from model import *

from domains.gridworld import *
from generators.obstacle_gen import *

import logging
import time
import math

def main(config,
         n_domains=3000,
         max_obs=30,
         max_obs_size=None,
         n_traj=1,
         n_actions=8,gen = False):
    # Correct vs total:
    logging.basicConfig(filename='./resources/logs/generalization/16_w_64_model',format='%(asctime)s-%(levelname)s:%(message)s', level=logging.INFO)
    correct, total = 0.0, 0.0
    # Automatic swith of GPU mode if available
    use_GPU = torch.cuda.is_available()
    # Instantiate a VIN model
    vin = VIN(config)
    # Load model parameters
    vin.load_state_dict(torch.load(config.weights))
    # Use GPU if available
    if use_GPU:
        vin = vin.cuda()
    counter,total_no_soln = 0,0
    global data
    data = []
    t_list = []
    total_dev_non_rel, total_dev_rel = 0.0,0.0
    total_dist, total_astar_dist = 0.0,0.0
    metrics = True #this enables displaying the distance left to reach goal upon a failure 
    dist_remain_avg = 0.0
    for dom in range(n_domains):
        if gen:
            print('Gen started')
            goal = [
            np.random.randint(config.imsize),
            np.random.randint(config.imsize)
            ]   
            obs = obstacles([config.imsize, config.imsize], goal, max_obs_size)
            # Add obstacles to map
            n_obs = obs.add_n_rand_obs(max_obs)
            # Add border to map
            border_res = obs.add_border()
            # Ensure we have valid map
            if n_obs == 0 or not border_res:
                continue
            start = None
        else:
            wpn = True
            # path = './resources/maps/'
            path = './resources/testing_maps/16x16/'
            mp, goal, start = open_map(dom,path)
            # path = './maps/8_data_300'
            # mp, goal, start = open_map_list(dom,path)
            mp[start[1]][start[0]] = 0 #Set the start position as freespace too
            mp[goal[1]][goal[0]] = 0 #Set the goal position as freespace too

            goal = [goal[1],goal[0]] #swap them around, for the row col format (x = col not row)
            start = [start[1],start[0]]
            obs = obstacles([config.imsize, config.imsize], goal, max_obs_size)
            obs.dom = mp
        
        # Get final map
        im = obs.get_final()


        #1 is obstacles. 
        #set obs.dom as the mp
        logging.debug('0 is obstacle ')
        logging.debug(' im: %s ', im)
        # Generate gridworld from obstacle map
        G = gridworld(im, goal[0], goal[1])
        # Get value prior
        value_prior = G.get_reward_prior()
        # Sample random trajectories to our goal
        states_xy, states_one_hot = sample_trajectory(G, n_traj,start,gen) #dijkstra trajectory 
        # print('states_xy', states_xy[0] , len(states_xy[0]))
        if gen and len(states_xy[0]) > 0:
            save_image(G.image,(goal[0],goal[1]),states_xy[0][0],states_xy, states_one_hot,counter) #this saves the maps 
        
        counter += 1
        t0 = time.time()
        for i in range(n_traj):
            if len(states_xy[i]) > 1:

                # Get number of steps to goal
                L = len(states_xy[i]) * 2
                # Allocate space for predicted steps
                pred_traj = np.zeros((L, 2))
                # Set starting position
                pred_traj[0, :] = states_xy[i][0, :]

                for j in range(1, L):
                    # Transform current state data
                    state_data = pred_traj[j - 1, :]
                    state_data = state_data.astype(np.int)
                    # Transform domain to Networks expected input shape
                    im_data = G.image.astype(np.int)
                    im_data = 1 - im_data
                    im_data = im_data.reshape(1, 1, config.imsize,
                                              config.imsize)
                    # Transfrom value prior to Networks expected input shape
                    value_data = value_prior.astype(np.int)
                    value_data = value_data.reshape(1, 1, config.imsize,
                                                    config.imsize)
                    # Get inputs as expected by network
                    X_in = torch.from_numpy(
                        np.append(im_data, value_data, axis=1)).float()
                    S1_in = torch.from_numpy(state_data[0].reshape(
                        [1, 1])).float()
                    S2_in = torch.from_numpy(state_data[1].reshape(
                        [1, 1])).float()
                    # Send Tensors to GPU if available
                    if use_GPU:
                        X_in = X_in.cuda()
                        S1_in = S1_in.cuda()
                        S2_in = S2_in.cuda()
                    # Wrap to autograd.Variable
                    X_in, S1_in, S2_in = Variable(X_in), Variable(
                        S1_in), Variable(S2_in)
                    # Forward pass in our neural net
                    _, predictions = vin(X_in, S1_in, S2_in, config)
                    _, indices = torch.max(predictions.cpu(), 1, keepdim=True)
                    a = indices.data.numpy()[0][0]
                    # Transform prediction to indices
                    s = G.map_ind_to_state(pred_traj[j - 1, 0],
                                           pred_traj[j - 1, 1])
                    ns = G.sample_next_state(s, a)
                    nr, nc = G.get_coords(ns)
                    pred_traj[j, 0] = nr
                    pred_traj[j, 1] = nc
                    if nr == goal[0] and nc == goal[1]:
                        # We hit goal so fill remaining steps
                        pred_traj[j + 1:, 0] = nr
                        pred_traj[j + 1:, 1] = nc
                        break
                # Plot optimal and predicted path (also start, end)
                if pred_traj[-1, 0] == goal[0] and pred_traj[-1, 1] == goal[1]:
                    logging.debug('#################### - Path Found map %s!\n', dom)
                    correct += 1
                    t1 = time.time()
                    t_list.append(t1-t0)
                    dev_rel,dev_non_rel,dist,astar_dist = deviation(states_xy[i],pred_traj,goal,total)
                    total_dev_rel += dev_rel
                    total_dev_non_rel += dev_non_rel
                    total_dist += dist
                    total_astar_dist += astar_dist
                    if config.plot == True:
                        visualize(G.image.T, states_xy[i], pred_traj)
                elif metrics:
                    d = dist_left(pred_traj,goal)
                    dist_remain_avg += d
                    if config.plot == True:
                        visualize(G.image.T, states_xy[i], pred_traj)
                total += 1



            elif wpn:
                total_no_soln += 1
        sys.stdout.write("\r" + str(int(
            (float(dom) / n_domains) * 100.0)) + "%")
        sys.stdout.flush()

    sys.stdout.write("\n")
    if total and correct:
        logging.info('Rollout Accuracy: %s',(100 * (correct / total)))
        logging.info('Rollout Accuracy Adjusted: %s',(100 * (correct / (total+total_no_soln))))
        logging.info('Total maps with no soln from Dijkstra %s', total_no_soln)
        logging.info('Total avg Rel Deviation %s', (total_dev_rel/total))
        logging.info('Total avg Non-Rel Deviation %s', (total_dev_non_rel/total))
        logging.info('Total avg VIN Distance %s', (total_dist/total))
        logging.info('Total avg Dijkstra Distance %s', (total_astar_dist/total))
        logging.info('Avg deviation from Dijkstra: %s', ((((total_astar_dist/total))-((total_dist/total)))/((total_astar_dist/total))))
        logging.info('Total elapsed time %s', (sum(t_list)/(total))) #TODO: Possibly add total no soln
        logging.info('Avg distance left when failed: %s ', (dist_remain_avg/(total-correct)) )
        logging.info('---------------------------------Done ------------------------------------')

    else:
        logging.info('No successes either vin or dijkstra')


def visualize(dom, states_xy, pred_traj):
    fig, ax = plt.subplots()
    implot = plt.imshow(dom, cmap="Greys_r")
    ax.plot(states_xy[:, 0], states_xy[:, 1], c='b', label='Optimal Path')
    ax.plot(
        pred_traj[:, 0], pred_traj[:, 1], '-X', c='r', label='Predicted Path')
    ax.plot(states_xy[0, 0], states_xy[0, 1], '-o', label='Start')
    ax.plot(states_xy[-1, 0], states_xy[-1, 1], '-s', label='Goal')
    legend = ax.legend(loc='upper right', shadow=False)
    for label in legend.get_texts():
        label.set_fontsize('x-small')  # the legend text size
    for label in legend.get_lines():
        label.set_linewidth(0.5)  # the legend line width
    plt.draw()
    plt.waitforbuttonpress(0)
    plt.close(fig)


def save_image(im, goal, start,states_xy,states_one_hot,counter):
    '''
    Saves the data made by generator as jsons. 
    '''
    s = config.imsize

    if len(states_xy[0]) == 0:

        im.tolist()[start_x][start_y] = 1
        start_xy = [0,0]
        mp = {
        'grid': im.tolist(),
        'goal': [goal[0],goal[1]],
        # 'start': int(start),
        'agent': start_xy}
        # 'states_xy': states_xy[0].tolist(),
        # 'states_one_hot': states_one_hot[0].tolist()
    else:
        mp = {
            'grid': im.tolist(),
            'goal': [goal[0],goal[1]],
            # 'start': int(start),
            'agent': states_xy[0][0].tolist()
            # 'states_xy': states_xy[0].tolist(),
            # 'states_one_hot': states_one_hot[0].tolist()   
    }
    data.append(mp)
    with open('./maps/' +str(s) + '_data_300' +  '.json', 'w') as outfile:
        json.dump(data,outfile)

def open_map(dom,path):
    '''
    Used to open a map json given dom and path, returns grid, goal and agent
    '''
    with open(str(path) + str(dom) +'.json') as json_file:
        data = json.load(json_file)
        logging.info('Opening file: ' + str(path) + str(dom) + '.json' )
        return data['grid'], data['goal'], data['agent']

def open_map_list(dom,path):
    with open(str(path) + '.json') as json_file:
        data = json.load(json_file)
        logging.info('Opening file: ' + str(path) + str(dom) + '.json' )
        return data[dom]['grid'], data[dom]['goal'], data[dom]['agent']

def deviation(optimal_path, pred_path,goal, map_num):
    optimal_path = np.array(optimal_path)
    optimal_path = 1.0 * optimal_path

    optimal_path_x = np.array(optimal_path[:,0])
    optimal_path_y = np.array(optimal_path[:,1])

    pred_path = np.unique(pred_path, axis=0) #removes duplicates at the end (when it reaches goal)

    #print('Shortened path' , pred_path)
    pred_path_x = np.array(pred_path[:,0])
    pred_path_y = np.array(pred_path[:,1])
    dist = 0.0
    astar_dist = 0.0
    prev = pred_path[0,:]
    total_diff_gen = 0
    for xy in pred_path[:,:]:

        diff = math.sqrt( ((1.0 * xy[0]- 1.0*prev[0])**2)+((1.0*xy[1] - 1.0*prev[1])**2))
        total_diff_gen += diff
        dist+= ((xy[0]-prev[0])**2 + (xy[1]-prev[1])**2)**0.5
        prev = xy 

    #prev = [0,0]
    #print('opt', optimal_path[0,:])
    prev = optimal_path[0,:]
    total_diff_optim = 0
    for xy in optimal_path[:,:]:
        # print('xy', xy)
        diff2 = math.sqrt( ((1.0 * xy[0]- 1.0*prev[0])**2)+((1.0*xy[1] - 1.0*prev[1])**2))
        total_diff_optim += diff2
        astar_dist+= ((xy[0]-prev[0])**2 + (xy[1]-prev[1])**2)**0.5
        prev = xy 
    
    dev_non_rel = abs(total_diff_optim-total_diff_gen)
    dev_rel = dev_non_rel/total_diff_optim #TODO: Add avg distance of gen trajectory
    return(dev_rel,dev_non_rel,dist,astar_dist)

def dist_left(pred_traj, goal):
    '''
    Finds the distance left between the point and the goal
    '''
    pred_traj = np.array(pred_traj) #euclidean distance or geometric distance ? use geometric
    x1,y1 = pred_traj[-1][0], pred_traj[-1][1]
    x2,y2 = goal[0],goal[1]
    dist = (((x2-x1)**2 + (y2-y1)**2))**0.5
    return dist        


if __name__ == '__main__':
    # Parsing training parameters
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--weights',
        type=str,
        default='trained/30k_no_block_dataset_vin_64x64.pth',
        help='Path to trained weights')
    parser.add_argument('--plot', action='store_true', default=False)
    parser.add_argument('--gen', action='store_true', default=False)
    parser.add_argument('--imsize', type=int, default=16, help='Size of image')
    parser.add_argument(
        '--k', type=int, default=20, help='Number of Value Iterations')
    parser.add_argument(
        '--l_i', type=int, default=2, help='Number of channels in input layer')
    parser.add_argument(
        '--l_h',
        type=int,
        default=150,
        help='Number of channels in first hidden layer')
    parser.add_argument(
        '--l_q',
        type=int,
        default=10,
        help='Number of channels in q layer (~actions) in VI-module')
    config = parser.parse_args()
    # Compute Paths generated by network and plot
    
    for i in range(1):
        main(config)
    # main(config)